Thank you for your continued help. I am using a ResNet50 with several heads built on top of it for multi task learning. It is actually a modification of this code. https://github.com/vita-epfl/openpifpaf

Also, did you make sure to write functions for the train and validation method?

Yes they are in separate functions. Line 89 of https://github.com/vita-epfl/openpifpaf/blob/master/openpifpaf/network/trainer.py
for epoch in range(start_epoch, epochs):
    self.train(train_scenes, epoch)

    self.write_model(epoch + 1, epoch == epochs - 1)
    self.val(val_scenes, epoch + 1)


Are you dealing with some dynamic sizes, e.g. varying batch shapes?

Kind of. I have 2 data loaders which have different batch sizes.
So my version has 2 dataloaders. One loads the mscoco dataset and another the volleyball dataset. It also has 2 heads and at each iteration, I pass the mscoco data through the ResNet basenet and through one of the heads before doing a backprop. Then I pass the volleyball data through ResNet and through the other head. Below is a simple description of my model.
Could this be the reason why I am running oom? Since the memory required by the model depends on the data that I am feeding through.
class Shell(torch.nn.Module):
    def __init__(self, base_net, head_nets):
        super(Shell, self).__init__()

        self.base_net    = base_net
        self.head_nets   = torch.nn.ModuleList(head_nets)
        self.head_pifpaf = self.head_nets[:2]
        self.head_crm    = self.head_nets[2:]

    def io_scales(self):
        return [self.base_net.input_output_scale // (2 ** getattr(h, '_quad', 0))
                for h in self.head_nets]

    def forward(self, x, head):  # pylint: disable=arguments-differ
        x = self.base_net(x)
        
        if head=="pifpaf":
            return [hn(x) for hn in self.head_pifpaf]
        if head=="crm":
            return [hn(x) for hn in self.head_crm]