I am kinda newbie and what I understand from your response for my case, input and ground truth should be modified as
train_ds = TensorDataset(x_train_, y_train)
train_dl = DataLoader(train_ds, batch_size=bs, shuffle=True)
train_dl = train_dl.to(dev)
valid_ds = TensorDataset(x_valid_, y_valid)
valid_dl = DataLoader(valid_ds, batch_size=bs * 2)
valid_dl = valid_dl.to(dev)
and for loss
loss=loss_func(model.feedforward(xb, sub),yb).to(dev)
But, it still does not work.