Fix

SINGROUP · Nov 27, 2023 · 1f15888 · 1f15888
1 parent 1db8682
commit 1f15888
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/integration_tests/test_train_posnet.py b/tests/integration_tests/test_train_posnet.py
@@ -39,7 +39,8 @@ def make_model(device, cfg):
         decoder_z_sizes=[5, 10, outsize],
         z_outs=[3, 3, 5, 8],
         peak_std=cfg["peak_std"],
-    ).to(device)
+        device=device
+    )
     criterion = nn.MSELoss(reduction="mean")
     optimizer = optim.Adam(model.parameters(), lr=cfg["lr"])
     lr_decay_rate = 1e-5