Merge pull request facebookresearch#68 from m-Just/patch-2

Update train.py
Marzsccc · Aug 25, 2021 · 4c16e32 · 4c16e32
2 parents 7f82a22 + 290d7d5
commit 4c16e32
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/domainbed/scripts/train.py b/domainbed/scripts/train.py
@@ -199,7 +199,7 @@ def save_checkpoint(filename):
             "model_num_classes": dataset.num_classes,
             "model_num_domains": len(dataset) - len(args.test_envs),
             "model_hparams": hparams,
-            "model_dict": algorithm.cpu().state_dict()
+            "model_dict": algorithm.state_dict()
         }
         torch.save(save_dict, os.path.join(args.output_dir, filename))