Memory error on some environments when using ensemble methods

Hi, I have an extremely strange bug that is slightly hard to reproduce :)
I am trying to integrate brax with SAC and want to use an ensemble version of the algorithm. When I use the ensemble, the brax environments (tested so far on halfcheetah, walker2d, grasp and hopper) raise the following error. The original SAC version doe not! In addition, the ant environment for some reason works as well.

I am wrapping the torch code in both cases with torch.no_grad() and detach from the tensor. The shapes going into brax are the same in either case. I think this is some strange memory allocation problem that happens when torch (internally?) duplicates the Tensor to compute the ensembled predictions, but that is pure conjecture.

Would be very happy about any hints.

```
Traceback (most recent call last):
File "sac/main.py", line 80, in main
next_state, reward, done, _ = env.step(action) # Step
File "sac/env/wrappers.py", line 104, in step
return self.env.step(action)
File "brax/envs/to_torch.py", line 59, in step
obs, reward, done, info = super().step(action)
File "lib/python3.10/site-packages/gym/core.py", line 289, in step
return self.env.step(action)
File "python3.10/site-packages/gym/wrappers/order_enforcing.py", line 11, in step
observation, reward, done, info = self.env.step(action)
File "python3.10/site-packages/brax/envs/wrappers.py", line 283, in step
self._state, obs, reward, done, info = self._step(self._state, action)
ValueError: INTERNAL: Address of buffer 11 must be a multiple of 10, but was 0x7f588c3d0848
````