Pytorch cuda tensor in module

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.weight = torch.nn.Parameter(torch.zeros(2, 1))
        self.bias = torch.nn.Parameter(torch.zeros(1))
        self.register_buffer('a_constant_tensor', torch.tensor([0.5]))

    def forward(self, x):
        # linear regression completely from scratch,
        # using parameters created in __init__
        x = torch.mm(x, self.weight) + self.bias + self.a_constant_tensor
        return x


model = Model().cuda()

Source: discuss.pytorch.org

Similar codes

71 votes384 views6 answers
Python ide online
Python
100 votes331 views6 answers
Pip3 upgrade
Python
3 votes365 views3 answers
Python find file name
Python
2 votes370 views2 answers
Pandas series to dataframe index as column
Python
11 votes339 views6 answers
Start python server
Python
3 votes334 views1 answer
Iteration
Python
0 vote360 views1 answer
Using python-poppler
Python
0 vote360 views1 answer
Progress apply
Python
0 vote358 views1 answer
Change every element of list python with map
Python
11 votes349 views6 answers
Pandas distinct
Python