Fixed it. Batch  and train function are now as follows:
def get_batches(arr_x, arr_y, batch_size):      
    # iterate through the arrays
    prv = 0
    for n in range(batch_size, arr_x.shape[0]+1, batch_size):
        x = arr_x[prv:n,:]
        y = arr_y[prv:n,:]
        prv = n
        yield x, y

model=Transformer(src_vocab_size=vocab_size, trg_vocab_size=vocab_size, src_pad_idx=0, trg_pad_idx=0).to(device)
opt=torch.optim.SGD(model.parameters(), lr=0.1)
loss_fn=nn.functional.cross_entropy
    
def train_model(src, trg, epochs=51, batch_size=32, classes=vocab_size):
    print(model)
    
    train_loss_list, validation_loss_list = [], []
    for e in range(epochs):
        model.train()
        
        N=np.ceil(src.shape[0]/batch_size)
        running_loss=0
        correct=0
        total=0
        
        if e%5==0: print("-"*25, f"Epoch {e + 1}","-"*25)
        for x, y in get_batches(src, trg, batch_size):
            x, y = torch.from_numpy(x.astype('int64')).to(device), torch.from_numpy(y.astype('int64')).to(device)
            opt.zero_grad()
            pred=model(x, y)
            predicted=torch.argmax(pred,dim=2)
            loss=loss_fn(pred, torch.nn.functional.one_hot(y, num_classes=classes).type(torch.FloatTensor))
            loss.backward()
            opt.step()
            if e%5==0:
                print(pred.shape)
                print(predicted[0])
                print(y[0])

            running_loss += loss.detach().item()
            train_loss_list += [running_loss]
            total+=y.size(0)*y.size(1)
            correct+=predicted.eq(y).sum().item()
            accu=100.*correct/total
            
        if e%5==0: print(f"Training loss: {running_loss/N:.4f}, Accuracy: {accu:.2f}%")
        print('...')
        
    return train_loss_list
train_loss_list = train_model(src=x4_int,trg=y4_int)