add cpu-only mode

2023-07-21 13:41:52 -07:00
parent 32aedaf9b6
commit ab0ce3c977
7 changed files with 39 additions and 42 deletions
--- a/train.py
+++ b/train.py
@@ -138,6 +138,8 @@ def run_test_eval(evaluator, model, dataloaders, writer, step):
            single_point=False,
            n_iters=6,
        )
+        if torch.cuda.is_available():
+            predictor.model = predictor.model.cuda()

        metrics = evaluator.evaluate_sequence(
            model=predictor,