The problem was solved by changing the forward function to:
def forward(self, x):
        x = self.base_model(x)
        x = self.relu(x)
        a = []
        for i in range(self.num_classes):
            a.append(self.heads[i](x[:, i*self.head_elements:i*self.head_elements+self.head_elements]).squeeze())
        
        return torch.stack(a, dim=1)

but I still don’t know why the former approach resulted in OOM error. The only difference that is made is that I store the output tensors in a python list instead of a new tensor. Is declaring a tensor in forward function or indexing it causing the problem?