xref: /aosp_15_r20/external/pytorch/test/bottleneck_test/test_cuda.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1# Owner(s): ["module: unknown"]
2
3import torch
4import torch.nn as nn
5
6
7class Model(nn.Module):
8    def __init__(self) -> None:
9        super().__init__()
10        self.linear = nn.Linear(20, 20)
11
12    def forward(self, input):
13        out = self.linear(input[:, 10:30])
14        return out.sum()
15
16
17def main():
18    data = torch.randn(10, 50).cuda()
19    model = Model().cuda()
20    optimizer = torch.optim.SGD(model.parameters(), lr=0.0001)
21    for i in range(10):
22        optimizer.zero_grad()
23        loss = model(data)
24        loss.backward()
25        optimizer.step()
26
27
28if __name__ == "__main__":
29    main()
30