How to set dtype for NN layers?

I have training data available as dtype = torch.float64 and I want to use this data type to train my model. Now it seems that the nn.Layer (e.g. nn.Linear) interface doesn’t offer the possibility of specifying a data type. Right now I get RuntimeError: Expected object of type torch.DoubleTensor but found type torch.FloatTensor because the dtypes are incompatible. How can I set the dtype of my NN layers to torch.float64? Do I have to assign the weights manually?

You can cast the whole layer to a new dtype using:

lin = nn.Linear(10, 10).double()
4 Likes

Thanks, that works! Is there a way to set it as the default for all layers or do I have to cast each layer manually? Something like a context manager would be handy; does it exist?

You can call it on the whole model which will cast all layers and parameters to the type.
Alternatively, you could set the default type at the beginning of your script using torch.set_default_dtype().

5 Likes

How can this be done in C++?

Also, in general, how can the .double() be achieved for a torch::Tensor object in C++?

Thank you.

You could use tensor = tensor.to(at::kDouble); to transform the tensor to an FP64 one and

AutoDefaultDtypeMode dtype_mode(default_dtype);

should work to define the default type.

2 Likes