Hi all,
I have an abstract class which inherits from nn.Module
and parent class which inherits from AbstractClass
when I try to assign the model to CUDA I don’t get any error but by looking at nvidia-smi the model is clearly not assigned. When I tried to use is_cuda
I got an error that my model does not recognize this method. I think the model is not inheriting the nn.Module
methods at all.
How can I fix this?
simple code to reproduce the problem:
class AbstractModel(nn.Module):
def __init__(self):
super().__init__()
self.max_objects_in_frame = 15
self.em_in_dim = 5
self.em_out_dim = 4
class Linear(AbstractModel):
def __init__(self, input_f, output_f):
super().__init__()
self.input_f = input_f
self.output_f = output_f
self.layer = nn.Linear(input_f, output_f)
def forward(self, x):
return self.layer(x)
if __name__ == '__main__':
model = Linear(10,20)
print(model.em_in_dim)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
print(model.is_cuda)