Add shape comments on every line that produces tensor

[kmkurn]

Something like:

# shape: (batch_size, seq_length)
words = words.transpose(0, 1).contiguous()

so it's easy to keep track of the shapes.