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

0 vote132 views1 answer
Matplotlib legend put larget dot
Python
6 votes156 views3 answers
How to split an input in python by comma
Python
3 votes213 views2 answers
Python iterate over string
Python
2 votes104 views1 answer
Python insert object into list
Python
0 vote178 views1 answer
Django model get field verbose name
Python Django
0 vote139 views1 answer
Plt force axis numbers
Python
1 vote137 views3 answers
Check space in string python
Python
0 vote150 views1 answer
_csv.Error: field larger than field limit (131072)
Python
0 vote148 views1 answer
Why python stops after plt.show()
Python
1 vote140 views1 answer
Django form list option
Python