Loss surface landscape by GPU

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)

Did you push the input data as well as the model to the GPU? If so, could you check if loss_landscapes creates new tensors in internal methods without specifying a device (the default would be CPU, which could create this issue)?

it solved without specifying a device

import loss_landscapes
import loss_landscapes.metrics

x, y = next(iter(testLoader))

# if haveCuda:
#   x,y = x.cuda(), y.cuda()

metric = loss_landscapes.metrics.Loss(criterion, x, y)

# compute loss data
loss_data = loss_landscapes.linear_interpolation(
    model_initial.cpu(), 
    model_final.cpu(), 
    metric, 
    40, 
    deepcopy_model=True
)