Fix mypy errors

HumanCompatibleAI · Jan 7, 2025 · 5820cfe · 5820cfe
1 parent f64b3fc
commit 5820cfe
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/imitation/algorithms/adversarial/common.py b/src/imitation/algorithms/adversarial/common.py
@@ -207,7 +207,7 @@ def __init__(
         self.debug_use_ground_truth = debug_use_ground_truth
         self.venv = venv
         self.gen_algo = gen_algo
-        self._reward_net = reward_net.to(gen_algo.device)
+        self._reward_net: reward_nets.RewardNet = reward_net.to(gen_algo.device)
         self._log_dir = util.parse_path(log_dir)
 
         # Create graph for optimising/recording stats on discriminator

diff --git a/src/imitation/policies/base.py b/src/imitation/policies/base.py
@@ -31,7 +31,7 @@ def _predict(
     ):
         np_actions = []
         if isinstance(obs, dict):
-            np_obs = types.DictObs(
+            np_obs: Union[types.DictObs, np.ndarray] = types.DictObs(
                 {k: v.detach().cpu().numpy() for k, v in obs.items()},
             )
         else: