What I know about the problem

Adam is stateful and requires a memory space proportional to the parameters in your model.
Model parameters must be loaded onto device 0
OOM occurs at state[‘exp_avg_sq’] = torch.zeros_like(p.data) which seems to be the last allocation of memory in the optimizer source code.
Neither manually allocating or use of nn.DataParallel prevents OOM error
Moved loss to forward function to reduce memory in training

Below are my training and forward methods
def train(dataloader, vocabulary_dict, epoch):
   
    model = ViralClassification(len(vocabulary_dict), 0.5, 6588 )#, device_ids=[0,1], output_device=1)
    model.to('cuda:0')
    model.print_gpu_memory_info()
    
    optimizer = optim.Adam(model.parameters(), lr=0.01)
    
    cudnn.benchmark = True
    cudnn.enabled = True


    for i in range(epoch):

        running_loss = 0.0

        for (i, (label, sequence)) in enumerate(dataloader):
            
            loss = model(sequence.to('cuda:0'), label.to('cuda:1'))#.to('cuda:0'))
            running_loss += loss.item()            
            loss.backward()
            del loss            
           
            print('gpu_memory_one in training_loop') 
            model.print_gpu_memory_info()
            torch.cuda.empty_cache()

            print('gpu_memory_two in training_loop')
            model.print_gpu_memory_info()
            
            print('bout to step')
            optimizer.step()
            
            
            optimizer.zero_grad() 
            print('bottom of training loop')

def forward(self, inputs, labels):

        inputs = self.embedding(inputs)#.to('cuda:0')
        #print('embedded completed')

        (inputs, hidden_state) = self.bilstm_layer(inputs)
        #print('bilstm completed')
        inputs.to('cuda:1')
        
        self.bilstm_layer.flatten_parameters()
        
        torch.cuda.synchronize('cuda:0')
        torch.cuda.synchronize('cuda:1') 
        
        inputs = self.attention_layer(inputs)
        #print('attention completed')

        inputs = inputs.view(-1, self.row*2*self.lstm_dim)#.to('cuda:1')
        #print('view transformation completed')
        
        inputs = self.mlp_one(inputs, self.relu_one)
        #print('mlp one completed')

        inputs = self.mlp_two(inputs, self.relu_two)
        #print('mlp two completed')

        torch.cuda.synchronize('cuda:0')
        torch.cuda.synchronize('cuda:1')
        
        logits = self._classify(inputs)
        #print('logits completed')

        torch.cuda.empty_cache()
        torch.cuda.synchronize('cuda:0')
        torch.cuda.synchronize('cuda:1')
        #self.print_gpu_memory_info() 

        loss = self.criterion(logits, labels)
        
        return loss

My OOM occurs when I perform optimizer.step.
My problem is that before optimizer.step my memory on device 1 has plenty of open room but since the optimizer performs it calculations on device 0, the OOM occurs.
Is this a problem that checkpointing may be able to solve?
Is it possible to change the location of the optimizer?