Fix numpy.float128 not found error (#270)

* Fix numpy.float128 not found error Fixes #268 * fix error * Fix typo: if->if not Co-authored-by: Adam Gleave <[email protected]>
HumanCompatibleAI · Mar 9, 2021 · 1c1bf24 · 1c1bf24
1 parent 9ddd6fa
commit 1c1bf24
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/imitation/data/types.py b/src/imitation/data/types.py
@@ -72,8 +72,8 @@ def _rews_validation(rews: np.ndarray, acts: np.ndarray):
             "rewards must be 1D array, one entry for each action: "
             f"{rews.shape} != ({len(acts)},)"
         )
-    if rews.dtype not in [np.float32, np.float64, np.float128]:
-        raise ValueError("rewards dtype {self.rews.dtype} not a float")
+    if not np.issubdtype(rews.dtype, np.floating):
+        raise ValueError(f"rewards dtype {rews.dtype} not a float")
 
 
 @dataclasses.dataclass(frozen=True)