Update test_generator.py

globus-labs · Nov 8, 2023 · 80e4f79 · 80e4f79
1 parent 9f9f7cb
commit 80e4f79
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/tests/test_generator.py b/tests/test_generator.py
@@ -2,10 +2,14 @@
 from pytest import fixture
 from mofa.generator import train_generator, run_generator
 import numpy as np
+import torch
+from 
 
 @fixture()
 def load_model():
-    ...
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    ddpm = DDPM.load_from_checkpoint(model, map_location=device).eval().to(device)
+    return ddpm
 
 def test_training():
     ...