I’m having a similar CUDA memory issue but I’m only running a shallow CNN…I was wondering if anyone could provide any insight as to where this memory leak is?
My model is pretty small:
class RegressionalNet(torch.nn.Module):
    
    def __init__(self):
        super().__init__()
        self.feature_extractor = torch.nn.Sequential(
            
            torch.nn.Conv2d(1,64,5,padding=2), 
            torch.nn.ReLU(),
            torch.nn.Conv2d(64,128,5,padding=2,bias=False),
            torch.nn.ReLU(),
            torch.nn.BatchNorm2d(128),
            torch.nn.MaxPool2d(2),
            torch.nn.Conv2d(128,256,5,padding=2,bias=False),
            torch.nn.ReLU(),
            torch.nn.BatchNorm2d(256),
            torch.nn.MaxPool2d(2))
 
        self.classifier = torch.nn.Sequential(
	        torch.nn.Dropout(0.5),
            torch.nn.Linear(256*16*16,256), # Fully connected layer 
            torch.nn.ReLU(),            
            torch.nn.Linear(256,1))
        
    def forward(self,x):
        features = self.feature_extractor(x)
        output = self.classifier(features.view(int(x.size()[0]),-1))
        return output

My training procedure looks as  follows:
def train(
          model: torch.nn.Module, 
          transforms, 
          data_path= DATA_PATH, 
          val_path= VAL_PATH, 
          num_epochs=101, 
          batch_size=64, 
          verbose=True,
          cube_length=640, img_size=(64, 64), 
          loss=torch.nn.MSELoss(), 
          lr_schedule=True, initial_lr=1e-3, suffix=""):

    data_path = os.path.abspath(data_path)
    val_path = os.path.abspath(val_path)	
    model = model.train()
    device = torch.device("cuda")
    model = model.to(device).to(torch.float)
    
    """ LOADING IN THE REAL TRAINING AND VALIDATION DATASETS """
    
    loader = DataLoader(FITSCubeDataset(data_path, cube_length, transforms, img_size), 
                        batch_size=batch_size, shuffle=True)                        
    validation_loader = DataLoader(FITSCubeDataset(val_path, cube_length, transforms, img_size), 
                                   batch_size=batch_size, shuffle=True)                                     
    
    optim = torch.optim.Adam(model.parameters(), initial_lr)	
    accuracies, val_accuracies, epochs, val_epochs = [0], [0], [0], [0]
	
    for i in range(num_epochs):
        print("Epoch %d of %d" % (i+1, num_epochs))
        _accuracies,_val_accuracies = [],[]    
        model.train(True) 
        for idx, (batch, target) in enumerate(tqdm(loader)):
            batch = batch.to(device).to(torch.float)
            if isinstance(loss, torch.nn.CrossEntropyLoss):
                target = target.to(device).to(torch.long)
            else:
                target = target.to(device).to(torch.float)
            pred = model(batch).reshape(-1)
            loss_value = loss(pred, target)
            optim.zero_grad()
            loss_value.backward()
            optim.step()
                
            ###Change the error metric here###

            _accuracies.append(loss_value)   
        epochs.append(i+1)                
        mean_accuracy = sum(_accuracies)/len(_accuracies)
        accuracies.append(mean_accuracy)
        print("Mean training loss: %f" % mean_accuracy)        

I’ve checked using watch -n 0.5 nvidia-smi to see if it really is a GPU memory issue and it does max out on epoch 1 when training on batch number of 606 every time.