SOM_LM-SNNs.py example breaks when running on a GPU

[ChatSam]

Runtime Error

 File "SOM_LM-SNNs.py", line 362, in <module>
    spikes=spike_record, assignments=assignments, n_labels=n_classes
  File "/home/ubuntu/bindsnet_temp/bindsnet/bindsnet/evaluation/evaluation.py", line 125, in all_activity
    rates[:, i] = torch.sum(spikes[:, indices], 1) / n_assigns
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

background info
python version - Python 3.6.12
pytorch version - torch 1.7.0

to reproduce the error, navigate to the examples/mnist folder, then run
python SOM_LM-SNNs.py --n_train 1000 --n_test 100 --gpu

[Hananel-Hazan]

Thanks for catching it!