I have the following model where encoder2 is a list of layers (not sequential). However, feeding the data on a GPU through encoder2 results in the data not being in GPU.
class model(nn.Module):
def __init__(self, args):
super(model, self).__init__()
self.encoder1 = make_mlp([6, 128], ["relu"])
self.encoder2 = [make_mlp([6, 128], ["relu"]) for _ in range(2)]
def forward(self, data, mode):
data = torch.stack([self.encoder2[i](data[i])] for i in range(2))
print(data.type())
def make_mlp(dim_list, activations, batch_norm=False, dropout=0):
if len(dim_list) == 0 and len(activations) == 0:
return nn.Identity()
assert len(dim_list) == len(activations)+1
layers = []
for dim_in, dim_out, activation in zip(dim_list[:-1], dim_list[1:], activations):
layers.append(nn.Linear(dim_in, dim_out))
if batch_norm:
layers.append(nn.BatchNorm1d(dim_out))
if activation == 'relu':
layers.append(nn.ReLU())
elif activation == 'leakyrelu':
layers.append(nn.LeakyReLU())
if dropout > 0:
layers.append(nn.Dropout(p=dropout))
return nn.Sequential(*layers)
How can I fix this?
net = model()
net = net.to('cuda:0')
data = data.cuda()
data = net(data)