Yes, the docs 1 mention the expected shape for each layer and I would stick it to it. My general rule is that a batch dimension is expected in nn.Modules. While e.g. linear layers could work with an input having a single dimension, you would have to verify what’s applied internally (which dimension is broadcasted etc.), so I would prefer to use the documented approach.