1import torch
2import torch.nn as nn
3
4class Model(nn.Module):
5 def __init__(self):
6 super(Model, self).__init__()
7 self.weight = torch.nn.Parameter(torch.zeros(2, 1))
8 self.bias = torch.nn.Parameter(torch.zeros(1))
9 self.register_buffer('a_constant_tensor', torch.tensor([0.5]))
10
11 def forward(self, x):
12 # linear regression completely from scratch,
13 # using parameters created in __init__
14 x = torch.mm(x, self.weight) + self.bias + self.a_constant_tensor
15 return x
16
17
18model = Model().cuda()
19