Solved by Mona_Jalal in post #2 
              
              
                never mind. I missed to convert one of the variables. I added     landmarks_batch = landmarks_batch.cpu() and now it works. 
# Helper function to show a batch
def show_landmarks_batch(sample_batched):
    """Show image with landmarks for a batch of samples."""
    images_batch, landmarks_batch = \
 …