Hello, I am trying to visualize my loss surface land-scape, I train the model on GPU colab
import loss_landscapes
import loss_landscapes.metrics
x, y = next(iter(testLoader))
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
x,y = x.to(device), y.to(device)
metric = loss_landscapes.metrics.Loss(criterion, x, y)
# compute loss data
loss_data = loss_landscapes.linear_interpolation(model_initial, model_final, metric, 40, deepcopy_model=True)
loss_data_fin = loss_landscapes.random_plane(model_final, metric,10, 40, normalization='filter', deepcopy_model=True)

RuntimeError                              Traceback (most recent call last)
<ipython-input-42-1a4435ce1a5c> in <module>()
----> 1 loss_data_fin = loss_landscapes.random_plane(model_final, metric,10, 40, normalization='filter', deepcopy_model=True)

1 frames
/usr/local/lib/python3.7/dist-packages/loss_landscapes/model_interface/model_parameters.py in filter_normalize_(self, ref_point, order)
    224             # normalize one-dimensional bias vectors
    225             if len(self.parameters[l].size()) == 1:
--> 226                 self.parameters[l] *= (ref_point.parameters[l].norm(order) / self.parameters[l].norm(order))
    227             # normalize two-dimensional weight vectors
    228             for f in range(len(self.parameters[l])):

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

it seem that it can’t work on GPU the last command:
oss_data_fin = loss_landscapes.random_plane(model_final, metric,10, 40, normalization='filter', deepcopy_model=True)