pytorch cuda tensor in module

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