Here is my loop. Thanks for helping a noob!
X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2, shuffle = True, random_state=0)
    
epochs = 1
LEARNING_RATE = 0.001
batch_size = 1000

model = NNTrain()
model.cuda()
optimizer = torch.optim.Adam(model.parameters(), lr = LEARNING_RATE)

model.train()
for i in range(epochs):
    for j in range(16):
        X_batch = torch.from_numpy(X_train[batch_size*j:batch_size*(j+1),:]).float().cuda()
        Y_batch = torch.from_numpy(Y_train[batch_size*j:batch_size*(j+1),:]).float().cuda()
        Y_batch_pred = model(X_batch)
        optimizer.zero_grad()
        loss = nn.L1Loss()(Y_pred, Y_train_tensor)
        loss.backward()
        optimizer.step()
    print(loss)