From d64df5dcf6fe1bc100fc932914aad53b1b3b4ab4 Mon Sep 17 00:00:00 2001
From: Yunfan Jiang <yunfanjyf@gmail.com>
Date: Tue, 18 Oct 2022 10:54:01 -0700
Subject: [PATCH] fix #2

---
 scripts/example.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scripts/example.py b/scripts/example.py
index 3022043..4966d9f 100644
--- 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(