training.py

for epoch in range(1, 101):
    running_loss = 0
    running_accuracy = 0
    running_loss_val = 0
    running_accuracy_val = 0
    
    start_time = time.time()
    # dataset class is assigned to dd variable
    dd.set_split('train')
    dataset = DataLoader(dd, batch_size=64, shuffle=True, collate_fn=padded)
    model.train()
    for batch_index, batch_dict in enumerate(dataset, 1):
        optimizer.zero_grad()

        x = batch_dict['x'].permute(1, 0)
        x = x.to(device)
        y = batch_dict['y'].to(device)

        output = model(x)[0]
        output = torch.softmax(output.squeeze(), dim=1)
        loss = criterion(output, y.type(torch.LongTensor).to(device))

        running_loss += (loss.item() - running_loss) / batch_index

        accuracy = compute_accuracy(y, output)
        running_accuracy += (accuracy - running_accuracy) / batch_index

        loss.backward()

        optimizer.step()
        
    dd.set_split('test')
    dataset = DataLoader(dd, batch_size=64, shuffle=True, collate_fn=padded)
    model.eval()        
    for batch_index, batch_dict in enumerate(dataset, 1):

        x = batch_dict['x'].permute(1, 0)
        x = x.to(device)
        y = batch_dict['y'].to(device)

        output = model(x)[0]
        output = torch.softmax(output.squeeze(), dim=1)
        loss = criterion(output, y.type(torch.LongTensor).to(device))

        running_loss_val += (loss.item() - running_loss_val) / batch_index

        accuracy = compute_accuracy(y, output)
        running_accuracy_val += (accuracy - running_accuracy_val) / batch_index