Merge pull request #11 from HumanCompatibleAI/typing_patch

Fix error from subscripting 'type' object

Author: PavelCz

src/reward_preprocessing/vis/reward_vis.py

@@ -1,7 +1,7 @@
 """Port of lucid.scratch.rl_util to PyTorch. APL2.0 licensed."""
 from functools import reduce
 import logging
-from typing import Optional
+from typing import List, Optional
 
 import lucent.optvis.param as param
 import lucent.optvis.render as render
@@ -15,7 +15,7 @@
 import reward_preprocessing.vis.objectives as objectives_rfi
 
 
-def argmax_nd(x: np.ndarray, axes: list[int], *, max_rep=np.inf, max_rep_strict=None):
+def argmax_nd(x: np.ndarray, axes: List[int], *, max_rep=np.inf, max_rep_strict=None):
     """Return the indices of the maximum value along the given axes.
 
     Args: