About example file&loss_landscape with GPU support

[mountains-high]

Hello!

Thank you for your fine work.
I wanted to reproduce the same experiments as shown in the example file.
But I got this error message while plotting loss_landscapes of original and pruned weights.
Could you help me please to use GPU? Thank you

 File "p.py", line 76, in <module>
    loss_data_fin = loss_landscapes.random_plane(model_final, metric, 1, STEPS, normalization='filter', deepcopy_model=True)
  File "/home/user/Desktop/loss_landscapes/main.py", line 285, in random_plane
    dir_one.filter_normalize_(start_point)
  File "/home/user/Desktop/loss_landscapes/model_interface/model_parameters.py", line 226, in filter_normalize_
    self.parameters[l] *= (ref_point.parameters[l].norm(order) / self.parameters[l].norm(order))
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!