Fixing binned network returning ndarray instead of tensor

Author: sash-a

.gitignore

@@ -9,5 +9,6 @@ bullet3/
 mujoco-py/
 ml-agents*/
 h-baselines/
+rllab/
 
 *mlruns*/

src/nn/nn.py

@@ -108,9 +108,9 @@ def __init__(self, layer_sizes: List[int], activation: nn.Module, env: gym.Env,
 
         super().__init__(layers, env.observation_space.shape, ob_clip)
 
-    def forward(self, inp: Tensor, **kwargs) -> np.ndarray:
+    def forward(self, inp: Tensor, **kwargs) -> Tensor:
         inp = clamp((inp - self._obmean) / self._obstd, min=-self.ob_clip, max=self.ob_clip)
-        a: np.ndarray = self.model(inp.float()).numpy()
+        a: Tensor = self.model(inp.float())
         ac_range = (self.ahigh - self.alow)[None, :]
 
         binned_ac = a.reshape((-1, self.adim, self.bins)).argmax(2)