one_iteration.py

%pip install torchviz==0.0.2

# one iteration
x = X[0]
y = Y[0]
y_predicted = model(x)
loss_tensor = loss(y_predicted, y)
loss_value = loss_tensor.data[0]
print(f"x: {x}, actual y: {y}, predicted y: {y}, loss: {loss_value}")
print(f"w: {w.data[0]}, b: {w.data[0]}")
loss_tensor.backward()
print(f"gradient for w is: {w.grad}, gradient for b is: {b.grad}")
with torch.no_grad():
  w -= w.grad * 1e-3
  b -= b.grad * 1e-3
print(f"updated w: {w.data[0]}, updated b: {w.data[0]}")
make_dot(loss_tensor)