I build a nn.Module that has a list containing some Linear.
I try to convert it to cuda but got error: RuntimeError: Expected object of backend CPU but got backend CUDA for argument #4 'mat1'
Is there any way to convert it? Here is some example code
class DeepLinear(nn.Module):
def __init__(self, dim: int, depth: int):
super(DeepLinear, self).__init__()
self.depth = depth
self.linear = []
for _ in range(depth):
self.linear.append(nn.Linear(dim, dim))
def forward(self, inputs):
for i in range(self.depth):
inputs = self.linear[i](inputs)
return inputs
a = DeepLinear(3, 5).cuda()
x = torch.ones(5, 3).cuda()
print(a(x)) # error