Describe the issue:

Hello!

While running one-shot NAS experiments on custom datasets, I encountered some inconsistencies in the usage of .to(device) and to_device().

EnasTrainer, on the one hand, uses the to_device() function from nni.retiarii.oneshot.pytorch.utils to transfer the tensor to the specified device. DartsTrainer and ProxylessTrainer, on the other hand, use the standard torch .to() function.

Using .to() is problematic, as the torch.utils.data.DataLoader instances might return batches that are not of type torch.Tensor and hence have no .to() method. The to_device function handles these cases.

Would it be possible to use to_device() across all Trainer classes?

Environment:

• NNI version: 2.3
• Python version: 3.8.8
• PyTorch version: 1.9.0

Thank you,
Thomas