I ran into the same issue. The solution described in HugginFace dataset error: RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.HalfTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor is what I am trying next. I’ll let you know if I run into more issues after upgrading.