fix: cuda device mismatch

vejvarm · Nov 17, 2023 · d0ac77a · d0ac77a
1 parent fa4b3e2
commit d0ac77a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.py b/train.py
@@ -218,7 +218,7 @@ def main():
         PREDICATE_POINTER: SingleTaskLoss,
         TYPE_POINTER: SingleTaskLoss,
         MULTITASK: MultiTaskLoss
-    }[args.task](ignore_index=vocabs[LOGICAL_FORM].stoi[PAD_TOKEN])
+    }[args.task](ignore_index=vocabs[LOGICAL_FORM].stoi[PAD_TOKEN], device=DEVICE)
 
     single_task_loss = SingleTaskLoss(ignore_index=vocabs[LOGICAL_FORM].stoi[PAD_TOKEN])