How to make predictions with a model that uses BatchNorm1d?

I’m writing a reinforcement learning program and I’m using nn.BatchNorm1d for the first time. The code trains fine but when it goes into evaluating the policy I get the “Expected more than 1 value per channel when training” error because nn.BatchNorm1d expects a batch of observations but I’m sending a single observation. I know why I’m getting the error, what I don’t know is how to set up the call to the policy to avoid the error, that is, to use the running estimates to normalize the single observation. Any suggestions?

You would have to call model.eval() to use the running stats in batchnorm layers and disable dropout layers. :wink:

Oh gee thanks. Did some search on the .eval() method. So it turns out that to evaluate we should

# Evaluation mode
model.eval()

with torch.no_grad():
    ...
    out_data = model(data)
    ...

and then turn the training method back on

# Training mode
...
model.train()
...

Again, first time me doing this.