# %% Импорт библиотек
import torch
import torch.nn as nn

# %% Входные переменные
input_ = 32*32
hidden_1 = 5
hidden_2 = 12
output = 2
learning_rate = 1e-3

weights_1 = torch.rand(size=(input_, hidden_1), requires_grad=True)
weights_2 = torch.rand(size=(hidden_1, hidden_2), requires_grad=True)
weights_3 = torch.rand(size=(hidden_2, output), requires_grad=True)

random_input = torch.rand(1, input_)

true_label = torch.zeros(1, output)
true_label[0][0] = 1
# %% Цикл обучения
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD([weights_1, weights_2, weights_3], lr=learning_rate)

hidden_layer_1 = random_input @ weights_1
hidden_layer_2 = hidden_layer_1 @ weights_2
output_layer = hidden_layer_2 @ weights_3

loss = criterion(output_layer, true_label)
print(loss)

optimizer.zero_grad()
loss.backward()
optimizer.step()

print(weights_1._grad)
print(weights_2._grad)
print(weights_3._grad)