Randomly get shape mismatch error

[vedantroy]

When doing inference, I randomly get the following error:

RuntimeError: The expanded size of the tensor (256) must match the existing size (287) at non-singleton dimension 0.  Target sizes: [256].  Tensor sizes: [287]