Merge branch 'main' of https://github.com/Robot-Learning-Library/VIMA …

…into main
zzx9636 · Nov 3, 2022 · dc76b08 · dc76b08
2 parents 824f80b + d64df5d
commit dc76b08
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/scripts/example.py b/scripts/example.py
@@ -81,7 +81,7 @@ def main(cfg):
     assert cfg.task in PARTITION_TO_SPECS["test"][cfg.partition]
 
     seed = 42
-    policy = create_policy_from_ckpt(cfg.ckpt, cfg.device)
+    policy = create_policy_from_ckpt(cfg.ckpt, cfg.device).to(cfg.device)
     env = TimeLimitWrapper(
         ResetFaultToleranceWrapper(
             make(