NaN values in Output when training randomly initialized transformer

ILIR: <p>Hello everyone, I am new to Pytorch and definitely not good, but I have to do this for class and am stuck at this problem. This is my code I am using to train a randomly initialized transformer.</p>
<p>import torch<br>
from torch import nn<br>
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler<br>
from transformers import get_linear_schedule_with_warmup<br>
from torch.optim import AdamW<br>
from torch.nn import TransformerEncoder, TransformerEncoderLayer<br>
import numpy as np</p>
<p>train_dataloader = DataLoader(train_dataset, sampler=RandomSampler(train_dataset), batch_size=16)<br>
validation_dataloader = DataLoader(validation_dataset, sampler=SequentialSampler(validation_dataset), batch_size=16)</p>
<p>class CustomTransformerEncoder(nn.Module):<br>
def <strong>init</strong>(self, num_tokens, dim_model, num_heads, num_encoder_layers, num_classes):<br>
super(CustomTransformerEncoder, self).<strong>init</strong>()<br>
self.embedding = nn.Embedding(num_tokens, dim_model)<br>
encoder_layer = TransformerEncoderLayer(d_model=dim_model, nhead=num_heads, batch_first=True)<br>
self.transformer_encoder = TransformerEncoder(encoder_layer=encoder_layer, num_layers=num_encoder_layers)<br>
self.final_layer = nn.Linear(dim_model, num_classes)</p>
<pre><code>def forward(self, src, mask=None):
    embedded_src = self.embedding(src)
    output = self.transformer_encoder(embedded_src, src_key_padding_mask=mask)
    output = self.final_layer(output[:, 0, :])
    return output
</code></pre>
<p>def init_weights(m):<br>
if isinstance(m, nn.Linear):<br>
nn.init.xavier_uniform_(m.weight)<br>
if m.bias is not None:<br>
nn.init.zeros_(m.bias)</p>
<p>model = CustomTransformerEncoder(<br>
num_tokens=30522,<br>
dim_model=512,<br>
num_heads=8,<br>
num_encoder_layers=6,<br>
num_classes=2,<br>
)</p>
<p>model.apply(init_weights)</p>
<p>device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)<br>
model.to(device)<br>
optimizer = AdamW(model.parameters(), lr=1e-4, eps=1e-8)<br>
epochs = 1<br>
total_steps = len(train_dataloader) * epochs<br>
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0, num_training_steps=total_steps)</p>
<p>loss_function = nn.CrossEntropyLoss()</p>
<p>def flat_accuracy(preds, labels):<br>
pred_flat = preds.argmax(dim=1).flatten()<br>
labels_flat = labels.flatten()<br>
return (pred_flat == labels_flat).type(torch.float).mean().item()</p>
<p>training_loss_values = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
validation_loss_values = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
validation_accuracy_values = <span class="chcklst-box fa fa-square-o fa-fw"></span></p>
<p>for epoch in range(epochs):<br>
print(‘======== Epoch {:} / {:} ========’.format(epoch + 1, epochs))<br>
total_train_loss = 0<br>
model.train()</p>
<pre><code>for step, batch in enumerate(train_dataloader):
    b_input_ids, b_input_mask, b_labels = tuple(t.to(device) for t in batch)

    optimizer.zero_grad()

    outputs = model(b_input_ids, b_input_mask.bool())

    loss = loss_function(outputs, b_labels)

    if torch.isnan(loss):
        print("NaN detected in loss.")
        break

    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()
    scheduler.step()

    total_train_loss += loss.item()

avg_train_loss = total_train_loss / len(train_dataloader)
print("  Average training loss: {0:.2f}".format(avg_train_loss))
training_loss_values.append(avg_train_loss)

print("Running Validation...")
model.eval()
total_eval_accuracy = 0
total_eval_loss = 0

for batch in validation_dataloader:
    b_input_ids, b_input_mask, b_labels = tuple(t.to(device) for t in batch)

    with torch.no_grad():
        outputs = model(b_input_ids, b_input_mask.bool())

        loss = loss_function(outputs, b_labels)
    total_eval_loss += loss.item()
    total_eval_accuracy += flat_accuracy(outputs, b_labels)

avg_val_accuracy = total_eval_accuracy / len(validation_dataloader)
print("  Validation Accuracy: {0:.2f}".format(avg_val_accuracy))
validation_accuracy_values.append(avg_val_accuracy)

avg_val_loss = total_eval_loss / len(validation_dataloader)
print("  Validation loss: {0:.2f}".format(avg_val_loss))
validation_loss_values.append(avg_val_loss)
</code></pre>
<p>and this is the output i get</p>
<p>======== Epoch 1 / 1 ========<br>
NaN detected in loss.<br>
Average training loss: 0.00<br>
Running Validation…<br>
Validation Accuracy: 0.50<br>
Validation loss: nan</p>
<p>how can I handle these NaN values?<br>
Thanks in advance!</p>