Allow `torch.device` Input for Device