test pipeline fix

KyloRen1 · Mar 13, 2023 · af7d839 · af7d839
1 parent f9a41ad
commit af7d839
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/configs/mobilenet_lstm_m2o.yaml b/configs/mobilenet_lstm_m2o.yaml
@@ -5,7 +5,7 @@ model_kwargs:
     pretrained: True
   temporal:
     name: lstm
-    hidden_size: 384
+    hidden_size: 256
     dropout: 0.3
     n_layers: 2
 data_kwargs:

diff --git a/test.py b/test.py
@@ -57,9 +57,11 @@ def parse_arguments():
     )
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
     model = create_model(experiment_config).to(device)
-    model.load_state_dict(torch.load(args.checkpoint_path))
+    pretrained_weights = torch.load(args.checkpoint_path, map_location=device)
+    model.load_state_dict(pretrained_weights)
+    print('loaded weights')
+
 
     test_model(
         experiment_config,