Error when training

[romilly]

When I run train.py, I get the following error:

  File "/home/romilly/git/external/VardaGPT/src/models/gpt2_associative.py", line 81, in forward
    batch_size, seq_len, _ = input_vectors.shape
ValueError: not enough values to unpack (expected 3, got 2)