Hello @albanD.  I have something similar.  My data class returns the image, label and ID.  How would I average over the ID.  each ID could have 1:N pictures in the models.  Code so far:

class SuperEncoder(nn.Module):
def init(self):
super(MyModel, self).init()
self.roofEncoder = nn.Sequential(
nn.Conv2d(3, 6, 3, 1, 1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(6, 12, 3, 1, 1),
nn.ReLU(),
nn.MaxPool2d(2)
)
    self.dwellingEconder = nn.Sequential(
        nn.Conv2d(1, 6, 3, 1, 1),
        nn.ReLU(),
        nn.MaxPool2d(2),
        nn.Conv2d(6, 12, 3, 1, 1),
        nn.ReLU(),
        nn.MaxPool2d(2)
    )
    
    self.fc1 = nn.Linear(54*54*16, 1000)
    
    self.fc2 = nn.Linear(54*54*16, 1000)
    
    
    
    self.fc_out(x)
        
    
def forward(self, x1, x2):
    x1 = self.roofEncoder(x1)
    x1 = x1.view(x1.size(0), -1)
    x1 = F.relu(self.fc1(x1))
    
    
    x2 = self.dwellingEncoder(x2)
    x2 = x2.view(x2.size(0), -1)
    x2 = F.relu(self.fc2(x2))


For more context see: How do I average photo feature outputs for later concatenation? 7