RuntimeError: Caught RuntimeError in replica x on device x. Original Traceback: ... RuntimeError: shape '[x, x, x]' is invalid for input of size xxx

Question

I am encountering a RuntimeError with the following message:

RuntimeError: Caught RuntimeError in replica 1 on device 1.
Original Traceback (most recent call last):
  File "path/to/file.py", line xx, in function_name
    padded_hidden, _ = pad_packed_sequence(packed_hidden, batch_first=True)
  File "/path/to/torch/nn/utils/rnn.py", line 314, in pad_packed_sequence
    padded_output, lengths = _VF._pad_packed_sequence(packed_sequence, batch_first, padding_value, total_length)
RuntimeError: shape '[56, 1, 768]' is invalid for input of size 41472

This error occurs in the context of a PyTorch Lightning training loop that is using DataParallel to train a model on multiple GPUs. The error seems to be related to padding packed sequences, as indicated by the line padded_hidden, _ = pad_packed_sequence(packed_hidden, batch_first=True).

What could be causing this error? How can I fix it?