PyTorch tensor basics and automatic differentiation

                  import torch

device = 'cuda' if torch.cuda.is_available() else 'cpu'
features = torch.tensor([[1.0, 2.0], [3.0, 4.0]], requires_grad=True, device=device)
weights = torch.tensor([[0.2], [0.8]], requires_grad=True, device=device)

logits = features @ weights
loss = (logits ** 2).mean()
loss.backward()

print('loss:', loss.item())
print('grad features:', features.grad)
print('grad weights:', weights.grad)