DirectML and CrossEntropyLoss

Hey there, I’ve been trying to make my own first model and I eventually realized I’d need to do a multioutput model.

I’ve done some tinkering and came to this train function but when trying to use directML I’ve ran into this error code:
RuntimeError: tensor.device().type() == at::DeviceType::PrivateUse1 INTERNAL ASSERT FAILED at “/home/vsts/work/1/s/pytorch-directml-plugin/torch_directml/csrc/dml/DMLTensor.cpp”:30, please report a bug to PyTorch. unbox expects Dml at::Tensor as inputs

I was wondering if there’s other people in my situation and whether I’ve missed something very simple