Solved by ptrblck in post #2 
              
              
                The nn.ReLU layer won’t apply any reduction to the data, but would apply the relu activation on each element as seen here: 
batch_size = 2
x = torch.randn(batch_size, 4096, 6)
relu = nn.ReLU()
out = relu(x)
print(out.shape)
> torch.Size([2, 4096, 6])

Depending on what dim1 represents you could redu…