From 96393c41c55c0a905eff035875b97f3b837225fe Mon Sep 17 00:00:00 2001
From: Lukas Schneider <32388635+e3ntity@users.noreply.github.com>
Date: Wed, 3 Apr 2024 13:25:58 -0700
Subject: [PATCH] Update runner.py to load models to correct device

---
 rsl_rl/runners/runner.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/rsl_rl/runners/runner.py b/rsl_rl/runners/runner.py
index b461a88..816f304 100644
--- a/rsl_rl/runners/runner.py
+++ b/rsl_rl/runners/runner.py
@@ -365,7 +365,7 @@ def _update(self) -> None:
 
     def load(self, path: str) -> Any:
         """Restores the agent and runner state from a file."""
-        content = torch.load(path)
+        content = torch.load(path, map_location=self.device)
 
         assert "agent" in content
         assert "data" in content