My workflow is kind of complex and I do not have a self contained example. But, I will try to explain it in as much detail as possible.I have a very large dataset that I can not load into memory. So I wrote a custom dataset class
class BigDataset(torch.utils.data.Dataset):
    #def __init__(self, data_paths, target_paths):
    def __init__(self, data_paths):
        self.data_memmaps = [np.load(path, mmap_mode='r') for path in data_paths]
        #self.target_memmaps = [np.load(path, mmap_mode='r') for path in target_paths]
        self.start_indices = [0] * len(data_paths)
        self.data_count = 0
        for index, memmap in enumerate(self.data_memmaps):
            self.start_indices[index] = self.data_count
            self.data_count += memmap.shape[0]

    def __len__(self):
        return self.data_count

    def __getitem__(self, index):
        memmap_index = bisect(self.start_indices, index) - 1
        index_in_memmap = index - self.start_indices[memmap_index]
        data = self.data_memmaps[memmap_index][index_in_memmap]
        return index, torch.from_numpy(data)

Next, I read the locations of all the files (my data is separated over multiple files)
    data_paths = [os.path.join(file_path, f'data/feature{index}.npy')
                  for index in range(2)]

    dataset = BigDataset(data_paths)

Since this dataset has both the train and validation data, I need to split it. Thus, I generate train and val indices and use the following code for train and val dataloader
    if args.distributed:
        #train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
        train_sampler = torch.utils.data.distributed.DistributedSampler(torch.utils.data.Subset(dataset, train_indices))
        val_sampler = torch.utils.data.distributed.DistributedSampler(torch.utils.data.Subset(dataset, val_indices))
    else:
        train_sampler = SubsetRandomSampler(train_indices)
        val_sampler = SubsetRandomSampler(val_indices)

    train_loader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size,
                                               num_workers=args.workers, sampler=train_sampler,
                                               pin_memory=True)
    val_loader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size,
                                             num_workers=args.worker, sampler=val_sampler,
                                             pin_memory=True)

I am pretty sure that the dataloader is causing the issue
Epoch: [0][ 0/51]       Time 220.202 (220.202)  Data 205.658 (205.658)  Loss 39.61639 (39.61639)        Accuracy   0.00 (  0.00)
Epoch: [0][ 0/51]       Time 220.181 (220.181)  Data 205.639 (205.639)  Loss 43.61139 (43.61139)        Accuracy   0.00 (  0.00)
Epoch: [0][ 0/51]       Time 220.229 (220.229)  Data 205.687 (205.687)  Loss 35.34707 (35.34707)        Accuracy   0.00 (  0.00)
Epoch: [0][ 0/51]       Time 220.228 (220.228)  Data 205.683 (205.683)  Loss 56.56057 (56.56057)        Accuracy   0.00 (  0.00)

Epoch: [0][ 1/51]       Time  0.917 (110.549)   Data  0.000 (102.820)   Loss 20.94585 (32.27862)        Accuracy   0.00 (  0.00)
Epoch: [0][ 1/51]       Time  0.917 (110.560)   Data  0.000 (102.829)   Loss 63.88563 (51.75101)        Accuracy   0.00 (  0.00)
Epoch: [0][ 1/51]       Time  0.917 (110.573)   Data  0.000 (102.844)   Loss 23.30010 (29.32359)        Accuracy   0.00 (  0.00)
Epoch: [0][ 1/51]       Time  0.917 (110.572)   Data  0.000 (102.842)   Loss 33.03528 (44.79793)        Accuracy   0.00 (  0.00)

I followed the same procedure described [here].(https://github.com/pytorch/examples/blob/master/imagenet/main.py 3)
Is the way I am loading my data (not directly in the memory) causing this issue?