jurgisp / pydreamer

PyTorch implementation of DreamerV2 model-based RL algorithm
MIT License
209 stars 48 forks source link

Batchnorm: Expected more than 1 value per channel #7

Open truncs opened 1 year ago

truncs commented 1 year ago

When the generator is trying run a NetworkPolicy it has a batch size of 1 which doesn't work well with BatchNorm Layers (see error in title). The error is due to the fact that the mean and var of the BatchNorm are still getting updated even with torch.no_grad and the error for batch size 1 is probably to indicate that this is not the intended behavior.

Do you think it is okay to run the NetworkPolicy in eval mode?

jurgisp commented 1 year ago

Good point, I agree the network should be run in eval mode in NetworkPolicy! I probably missed it, because I did not use layers that need eval setting during inference.