GroupLinearLayer should add "device" parameter

[ildefons]

Hi RIM dev team,

Code fails when device = 'cuda'
It can be easily solved adding an extra parameter "device" to all "Group" classes.

Thank you for the great RIM implementation,
Ildefons

[dido1998]

HI @ildefons, could you share the error message with device='cuda'?

[ildefons]

Error message:

RuntimeError Traceback (most recent call last)
in
1 for x in xs:
2 print(1)
----> 3 hs, cs = rim_model(x, hs, cs)

~\anaconda3\envs\eg2\lib\site-packages\torch\nn\modules\module.py in call(self, *input, **kwargs)
548 result = self._slow_forward(*input, **kwargs)
549 else:
--> 550 result = self.forward(*input, **kwargs)
551 for hook in self._forward_hooks.values():
552 hook_result = hook(self, input, result)

~\OneDrive\Documentos\YK\eg\Recurrent-Independent-Mechanisms\RIM.py in forward(self, x, hs, cs)
249
250 # Compute input attention
--> 251 inputs, mask = self.input_attention_mask(x, hs)
252 h_old = hs * 1.0
253 if cs is not None:

~\OneDrive\Documentos\YK\eg\Recurrent-Independent-Mechanisms\RIM.py in input_attention_mask(self, x, h)
177 key_layer = self.key(x)
178 value_layer = self.value(x)
--> 179 query_layer = self.query(h)
180
181 key_layer = self.transpose_for_scores(key_layer, self.num_input_heads, self.input_key_size)

~\anaconda3\envs\eg2\lib\site-packages\torch\nn\modules\module.py in call(self, *input, **kwargs)
548 result = self._slow_forward(*input, **kwargs)
549 else:
--> 550 result = self.forward(*input, **kwargs)
551 for hook in self._forward_hooks.values():
552 hook_result = hook(self, input, result)

~\OneDrive\Documentos\YK\eg\Recurrent-Independent-Mechanisms\RIM.py in forward(self, x)
31 x = x.permute(1,0,2)
32
---> 33 x = torch.bmm(x,self.w)
34 return x.permute(1,0,2)
35

RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 'mat2' in call to _th_bmm