error while using Recurrent block in lava-dl

[franzhd]

Describe the bug I'm trying to train recurrent neurons, but i'm not a ble due to some errors in the code. i'm posting down the code of the net itself. To reproduce current behavior Steps to reproduce the behavior:

1. When I run this code (add code or minimum test case) ...

   
   class Network(torch.nn.Module):
   def __init__(self, input, output, treshold, voltage_decay):
       super(Network, self).__init__()
   
       neuron_params = {
               'threshold'     : treshold,
               'current_decay' : 1,               
               'voltage_decay' : voltage_decay,
               'tau_grad'      : 1,
               'scale_grad'    : 1
           }
       neuron_params_drop = {**neuron_params, 'dropout' : slayer.neuron.Dropout(p=0.1),}
   
       self.blocks = torch.nn.ModuleList([
   
               slayer.block.cuba.Dense(neuron_params_drop, input, 128, weight_norm=True, delay=True),
               slayer.block.cuba.Recurrent(neuron_params_drop, 128,256,weight_norm=True,  delay=False),
               slayer.block.cuba.Recurrent(neuron_params_drop, 256,256, weight_norm=True, delay=False),
               slayer.block.cuba.Dense(neuron_params, 256, 128,  weight_norm=True, delay=True),
               slayer.block.cuba.Dense(neuron_params, 128, output)
           ])
   
   def forward(self, x):
       count = []
       event_cost = 0
   
           # forward computation is as simple as calling the blocks in a loop
       x = self.blocks[0](x)
       x = self.blocks[1](x)
       x = self.blocks[2](x)
       x = self.blocks[3](x)
       x = self.blocks[4](x)
           # if hasattr(block, 'neuron'):
           #     event_cost += event_rate_loss(x)
           #     count.append(torch.sum(torch.abs((x[..., 1:]) > 0).to(x.dtype)).item())
   
       return x #, event_cost , torch.FloatTensor(count).reshape((1, -1)).to(x.device)
   
   def grad_flow(self, path):
       # helps monitor the gradient flow
       grad = [b.synapse.grad_norm for b in self.blocks if hasattr(b, 'synapse')]
   
       plt.figure()
       plt.semilogy(grad)
       plt.savefig(path + 'gradFlow.png')
       plt.close()
   
       return grad
   
   def export_hdf5(self, filename):
       # network export to hdf5 format
       h = h5py.File(filename, 'w')
       layer = h.create_group('layer')
       for i, b in enumerate(self.blocks):
           b.export_hdf5(layer.create_group(f'{i}'))

2. I get this error ...

Traceback (most recent call last): File "/root/lava-dl_experiment/src/experiment13.py", line 194, in main() File "/root/lava-dl_experiment/src/experiment13.py", line 152, in main output = assistant.train(input.to(device), label) File "/opt/conda/lib/python3.10/site-packages/lava/lib/dl/slayer/utils/assistant.py", line 121, in train output = self.net(input) File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl return forward_call(*args, kwargs) File "/root/lava-dl_experiment/src/experiment13.py", line 56, in forward x = self.blocks1 File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl return forward_call(*args, *kwargs) File "/opt/conda/lib/python3.10/site-packages/lava/lib/dl/slayer/block/base.py", line 1476, in forward x = recurrent.custom_recurrent(z, self.neuron, self.recurrent_synapse) File "/opt/conda/lib/python3.10/site-packages/lava/lib/dl/slayer/utils/recurrent.py", line 42, in custom_recurrent return CustomRecurrent.apply(z, neuron, recurrent_mat) File "/opt/conda/lib/python3.10/site-packages/torch/autograd/function.py", line 506, in apply return super().apply(args, kwargs) # type: ignore[misc] File "/opt/conda/lib/python3.10/site-packages/lava/lib/dl/slayer/utils/recurrent.py", line 64, in forward feedback = torch.matmul(spike[..., 0], recurrent_mat_T) RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument mat2 in method wrapper_CUDA_mm)

**Expected behavior**

**Screenshots**
If applicable, add screenshots to help explain your problem. Remove section otherwise.

**Environment (please complete the following information):**
 - docker https://github.com/franzhd/lava-dl_docker
 - Lava version  0.9.0
 - lava-dl version 0.5.0

**Additional context**
Add any other context about the problem here. Remove section otherwise.