Modules should be created on the default device

[brianberns]

The following F# program crashes:

open TorchSharp
open type torch
open type torch.nn

torch.set_default_device(torch.CUDA)

let embedding = Embedding(2, 3)
let tensor = tensor([|1|])
tensor --> embedding   // boom

The error message is: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0!. This occurs because the embedding is created on the CPU, even though CUDA is set as the default device.

---

It looks like all modules created in this way ignore the default device, although this doesn't cause a crash for other module types. I'm not sure if this lack of a crash is the expected behavior or not, but it surprised me. For example, the following code works, even though the linear module is created on the CPU:

torch.set_default_device(torch.CUDA)

let linear = Linear(2, 3)
let tensor = tensor([|1.0f; 2.0f|])
tensor --> linear

The output tensor in this case is on the CUDA device, even though the linear module is running on the CPU.

[ds5678]

The Linear module is implemented in managed code, which is why it works.

https://github.com/dotnet/TorchSharp/blob/3760ba3e97ca09a35ad8eab1d5128b62d6d2c2ef/src/TorchSharp/NN/Linear.cs

Embedding is handled by the native torch library.

TorchSharp/src/TorchSharp/NN/Embedding.cs

Lines 60 to 68 in 3760ba3

	public static Embedding Embedding(long num_embeddings, long embedding_dims, long? padding_idx = null, double? max_norm = null, double norm_type = 2.0, bool scale_grad_by_freq = false, bool sparse = false, Device? device = null, ScalarType? dtype = null)
	{
	var res = THSNN_Embedding_ctor(num_embeddings, embedding_dims,
	padding_idx.HasValue ? padding_idx.Value : -1, padding_idx.HasValue,
	max_norm.HasValue ? max_norm.Value : 0.0, max_norm.HasValue,
	norm_type, scale_grad_by_freq, sparse, out var boxedHandle);
	if (res == IntPtr.Zero) { torch.CheckForErrors(); }
	return new Embedding(res, boxedHandle).MoveModule<Embedding>(device,dtype);
	}

As for moving to the default device, this seems to be a PyTorch "bug" which is more of a feature at this point. I doubt it will be changed there.

I too have encountered the annoyance of having to move my modules over to the default device. I would be in favor of having TorchSharp do this automatically. I think that a modification to this function would be all that's required.

TorchSharp/src/TorchSharp/NN/Module.cs

Lines 1138 to 1145 in 3760ba3

	internal T MoveModule<T>(Device? device, ScalarType? dtype) where T : Module
	{
	T module = (T)this;
	return device != null ?
	(dtype.HasValue ? (T)module._to(device, dtype.Value, false) : (T)module._to(device.type, device.index, false)) :
	(dtype.HasValue ? (T)module._to(dtype.Value, false) : module);
	}

[ozanMSFT]

@brianberns , thanks for bringing this up.

I've created a code to reproduce similar scenario with pytorch and it worked as expected (by considering CUDA as device for both Embedding and tensor)

import torch
import torch.nn as nn

# Set the default device to CUDA
torch.set_default_device('cuda')

embedding = nn.Embedding(2, 3)

tensor = torch.tensor([1])

result = embedding(tensor)

print(result)

---

It seems there are some missing checks (or wrong default assumptions) while creation of some modules.

We'll investigate and try to address this issue to give ability to TorchSharp for considering default device automatically.