and this is the rest of my code:
import numpy as np
from tqdm import tqdm
import torch.optim as optim
training_data = np.load("/content/thirdy/training_data_grayscale.npy", allow_pickle = True)
x = torch.Tensor([i[0] for i in training_data]).view(-1, 100, 100)
y = torch.Tensor([i[1] for i in training_data])
VAL_PCT = 0.1
val_size = int(len(x)*VAL_PCT)
print(val_size)
train_x = x[:-val_size]
train_y = y[:-val_size]
test_x = x[-val_size:]
test_y = y[-val_size:]
import torch.optim as optim
optimizer = optim.Adam(model.parameters(), lr = 0.001)
loss_function = nn.CrossEntropyLoss()
BATCH_SIZE = 1
EPOCHS = 1
for epoch in range(EPOCHS):
for i in tqdm(range(0, len(train_x), BATCH_SIZE)): 

    batch_x = train_x[i:i+BATCH_SIZE].view(-1, 1, 50, 50)

    print(batch_x.shape)

    batch_y = train_y[i:i+BATCH_SIZE]

    print(batch_y.shape)

    model.zero_grad()

    outputs = model(batch_x)

    loss = loss_function(outputs, batch_y)

    loss.backward()

    optimizer.step()    # Does the update

print(f"Epoch: {epoch}. Loss: {loss}")