Update test_metric_lightning.py

Lightning-AI · Nov 3, 2020 · d8ff339 · d8ff339
1 parent 76de243
commit d8ff339
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/metrics/test_metric_lightning.py b/tests/metrics/test_metric_lightning.py
@@ -120,7 +120,7 @@ def training_step(self, batch, batch_idx):
 
  # save to file and re-load to ensure export + load still works for inference
  path = os.path.join(tmpdir, "tmp_script.pt")
- torch.jit.save(script_model, path)
+ script_model.save(path)
  load_script_model = torch.jit.load(path)
  load_script_model_output = load_script_model(rand_input)
  assert torch.allclose(output, load_script_model_output)