test/bottleneck_test/test.py

import torch

x = torch.ones((3, 3), requires_grad=True)
(3 * x).sum().backward()