Sorry for the late reply.
I agree with you in case of using by-default torch.nn layers like 2d convolution, which expects to have an extra dimension for the batch. I was wondering what about custom nn.modules
Toy example here
import torch.nn as nn
import torch
import numpy as np

class trial(nn.Module):
    def __init__(self):
        super(trial, self).__init__()
        self.conv = nn.Conv2d(1, 4,kernel_size=3)
    def forward(self,input):
        dims = input.size()
        x = []
        for i in range(dims[0]):
            x.append(self.conv(input[i,:,:,:,:]))

        return dims, torch.stack(x)
class trial2(nn.Module):
    def __init__(self):
        super(trial2, self).__init__()
        self.conv = nn.Conv2d(1, 4,kernel_size=3)
    def forward(self,input):
        dims = input.size()
        x = self.conv(input)
        return x
x = torch.rand(10,1,50,50)
print 'X size {0}'.format(x.size())
m = trial()

y = m(x)
print 'Y size {0}'.format(y[0])

Imagine I have data of 4 dimensions + 1 for batch = 5D
So lets imagine i have groups of 10 images to be processed:
Batch x group  x channels x H x W
[5,2,3,50,50]
In this case I have some questions:
is using a for loop over dim 0 the only way to process this batch  ?
is the computational graph  the same doing m(input) for an input of size [5,2,3,50,50] than doing conv2d(input) for an input size [2x5,3,50,50]?  (This is, reshaping all the samples in one dimension)
I’ve seen that in case 2) the graph is

image.png1141×703 35.4 KB

Meanwhile in case of using for loop and more dimensions the graph is

image.png1473×498 39.5 KB

Well obviously it is not, what is the proper way of dealing with data of custom dimensionality not to get improper graphs? Are there some rules about how to proceed at the time of working with multimodal NN?