Sorry not module but the part when I am “training my model”, the GPU utilization shows only  5-6 % usage and keeps on processing. For 400k records, it disconnects finally and doesnt give any output while it works fine for small dataset. I have no idea as to why the GPU usage is only 5-6% while training my model. Is there something incorrect.
Training the model
model = Net().cuda()
optimizer = optim.Adam(model.parameters(), lr=0.001)

loss_func = nn.NLLLoss()

epochs = 3
loss_list = []
correct = 0
model.train()

for epoch in range(epochs):
    total_loss = []
    
    for i, (X_tr, Y_tr) in enumerate(train_ldr):
        
  # get the inputs
      X_tr = X_tr.unsqueeze(0).cuda()
     
      Y_tr = Y_tr.type(torch.LongTensor).cuda()
      
      optimizer.zero_grad()

        # Forward pass
      output = model(X_tr.unsqueeze(0)).cuda()
      
      pred = output.argmax(dim=1, keepdim=True)
#      correct += pred.eq(Y_tr.view_as(pred)).sum().item()
      loss = loss_func(output, Y_tr.type(torch.LongTensor).cuda())
        # Backward pass
      loss.backward()
        # Optimize the weights
      optimizer.step()
        
      total_loss.append(loss.item())
#      print(correct, Y_tr, pred)

      loss_list.append(sum(total_loss)/len(total_loss))
    print('Training [{:.0f}%]\tLoss: {:.4f}'.format(
        100. * (epoch + 1) / epochs, loss_list[-1]))
    ~~~