Loss, Logits error while training

[niranjanakella]

There is a assignment error in the train.py script where in the loss and logits are considered to be 'str' type after the assignment and hence have to be updated.

Line: 108 and 146

loss, logits = model(texts, attention_mask=masks, labels=labels)

Here the loss variable is assigned as a 'str' type hence the following loss.backward() would fail stating that a 'str' type doesn't have a backward method.

[niranjanakella]

This can be corrected by re-assigning the loss and logits to the corresponding model output values.

Before

loss, logits = model(texts, attention_mask=masks, labels=labels)

After

model_out = model(texts, attention_mask=masks, labels=labels)
loss, logits = model_out.loss, model_out.logits

NOTE: PR #54 has been raised