How to add new nn.Module with GPU & CPU support

1

HI, I meet a problem when adding a new NN Module. There are two LSTM in the new module and LSTM need initial hidden state h_0.

init_hidden = (Variable(torch.zeros(1, inputs.size(1), self.hiddenchannel)),
               Variable(torch.zeros(1, inputs.size(1), self.hiddenchannel)))

It worked when I run outputs = model(inputs) with CPU mode. However if I run model = model.cuda(), outputs = model(inputs) Then error occurs.

I modify h_0 as:

init_hidden = (Variable(torch.zeros(1, inputs.size(1), self.hiddenchannel).cuda()),
               Variable(torch.zeros(1, inputs.size(1), self.hiddenchannel).cuda()))

and run model = model.cuda(), outputs = model(inputs) It works correctly.

So how to init h_0 thus it works in CPU & GPU mode when running outputs = model(inputs) and model = model.cuda(), outputs = model(inputs)?

Thanks

2

@smth I would appreciate it if you can give some help on this problem. Thanks you .

3

Hi,

Depending on where you create your hidden state you have few possibilities.
If you have access to the current input, you can do input_tensor.new(1, inputs.size(1), self.hiddenchannel).zero_().
If you don’t, you can use .type_as(input_tensor) at the beginning of your forward function.