Hi, you are welcome. You don’t need to worry about the dimension, because PyTorch will consider the first dimension of the input to the network as BATCH dimension (as long as your network is inherited from nn.Module). For example, the input size of your network is 1xD, and you feed BxD arrays to your network, that will definitely work well, PyTorch will consider B as the BATCH dimension. And you don’t need to recode, the original code is ok. Just fix the batch update part. Hope that helps, Good luck!