Tensor is not conforming the model device

1

I added the following function to my model architecture to reshape my parameter vector L_vec inside the model to a triangular shape matrix. However I need to map manually the device of matrix to the device of L_vec.
I am not sure what could be the reason.

def train_tensor(self, L_vec, n): 
        tri_mat = torch.zeros((n, n)).to(L_vec.device) 
        tril_indices = torch.tril_indices(row=n, col=n, offset=0)
        tri_mat[tril_indices[0], tril_indices[1]] = L_vec.squeeze(1)
        return tri_mat
2

Do you see any errors and if so, could you post the error message please?
Or does L_vec.device not give the desired device?

3

I wonder why I need manually map the device of tri_mat to the device L_vec while they are located in the same model.

.

4

You are recreating tri_mat in the forward method, so this tensor is not part of the model but just a temporary tensor, so it won’t use any specific device.
Also note, that a module doesn’t necessarily belong to a single device, since you can use model sharding etc.