From 1f1588851837526c40b30e6ac3ad55457b3fb282 Mon Sep 17 00:00:00 2001 From: NikoOinonen Date: Mon, 27 Nov 2023 18:17:38 +0200 Subject: [PATCH] Fix --- tests/integration_tests/test_train_posnet.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/integration_tests/test_train_posnet.py b/tests/integration_tests/test_train_posnet.py index 593abe7..e0c2f97 100644 --- 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