Pytorch Crossentropy loss - RuntimeError: "nll_loss_forward_reduce_cuda_kernel_2d_index" not implemented for 'Float'

[animeshkumarpaul]

This below code is working fine:

loss = nn.CrossEntropyLoss()
input = torch.randn(3, 5, requires_grad=True)
target = torch.randn(3, 5).softmax(dim=1).float()
output = loss(input, target)

But, when I am using the cross-entropy loss inside the skorch, I am getting an error:

device='cuda:5'

net = NeuralNetClassifier(
    ClassifierModule,
    module__layer_sizes= [...,...,...],
    module__input_size=X.shape[1],
    max_epochs=200,
    batch_size=50000,
    device= device,
   
    criterion=nn.CrossEntropyLoss,
    criterion__reduction='mean',
     

    optimizer=torch.optim.Adam,
    lr=0.01,
    optimizer__weight_decay=1e-4
)

Error

----> 3 net.fit(X_train_NN,y_train)

File ~/anaconda3/envs/xgb3_ray2_pytorch/lib/python3.8/site-packages/skorch/classifier.py:143, in NeuralNetClassifier.fit(self, X, y, **fit_params)
    132 """See ``NeuralNet.fit``.
    133 
    134 In contrast to ``NeuralNet.fit``, ``y`` is non-optional to
   (...)
    138 
    139 """
    140 # pylint: disable=useless-super-delegation
    141 # this is actually a pylint bug:
    142 # https://github.com/PyCQA/pylint/issues/1085
--> 143 return super(NeuralNetClassifier, self).fit(X, y, **fit_params)

File ~/anaconda3/envs/xgb3_ray2_pytorch/lib/python3.8/site-packages/skorch/net.py:1302, in NeuralNet.fit(self, X, y, **fit_params)
   1299 if not self.warm_start or not self.initialized_:
   1300     self.initialize()
-> 1302 self.partial_fit(X, y, **fit_params)
   1303 return self

File ~/anaconda3/envs/xgb3_ray2_pytorch/lib/python3.8/site-packages/skorch/net.py:1261, in NeuralNet.partial_fit(self, X, y, classes, **fit_params)
   1259 self.notify('on_train_begin', X=X, y=y)
   1260 try:
-> 1261     self.fit_loop(X, y, **fit_params)
   1262 except KeyboardInterrupt:
   1263     pass

File ~/anaconda3/envs/xgb3_ray2_pytorch/lib/python3.8/site-packages/skorch/net.py:1173, in NeuralNet.fit_loop(self, X, y, epochs, **fit_params)
   1170 for _ in range(epochs):
   1171     self.notify('on_epoch_begin', **on_epoch_kwargs)
-> 1173     self.run_single_epoch(iterator_train, training=True, prefix="train",
   1174                           step_fn=self.train_step, **fit_params)
   1176     self.run_single_epoch(iterator_valid, training=False, prefix="valid",
   1177                           step_fn=self.validation_step, **fit_params)
   1179     self.notify("on_epoch_end", **on_epoch_kwargs)

File ~/anaconda3/envs/xgb3_ray2_pytorch/lib/python3.8/site-packages/skorch/net.py:1209, in NeuralNet.run_single_epoch(self, iterator, training, prefix, step_fn, **fit_params)
   1207 for batch in iterator:
   1208     self.notify("on_batch_begin", batch=batch, training=training)
-> 1209     step = step_fn(batch, **fit_params)
   1210     self.history.record_batch(prefix + "_loss", step["loss"].item())
   1211     batch_size = (get_len(batch[0]) if isinstance(batch, (tuple, list))
   1212                   else get_len(batch))

File ~/anaconda3/envs/xgb3_ray2_pytorch/lib/python3.8/site-packages/skorch/net.py:1088, in NeuralNet.train_step(self, batch, **fit_params)
   1080     self.notify(
   1081         'on_grad_computed',
   1082         named_parameters=TeeGenerator(self.get_all_learnable_params()),
   1083         batch=batch,
   1084         training=True,
   1085     )
   1086     return step['loss']
-> 1088 self._step_optimizer(step_fn)
   1089 return step_accumulator.get_step()

File ~/anaconda3/envs/xgb3_ray2_pytorch/lib/python3.8/site-packages/skorch/net.py:1043, in NeuralNet._step_optimizer(self, step_fn)
   1041     optimizer.step()
   1042 else:
-> 1043     optimizer.step(step_fn)

File ~/anaconda3/envs/xgb3_ray2_pytorch/lib/python3.8/site-packages/torch/optim/optimizer.py:280, in Optimizer.profile_hook_step.<locals>.wrapper(*args, **kwargs)
    276         else:
    277             raise RuntimeError(f"{func} must return None or a tuple of (new_args, new_kwargs),"
    278                                f"but got {result}.")
--> 280 out = func(*args, **kwargs)
    281 self._optimizer_step_code()
    283 # call optimizer step post hooks

File ~/anaconda3/envs/xgb3_ray2_pytorch/lib/python3.8/site-packages/torch/optim/optimizer.py:33, in _use_grad_for_differentiable.<locals>._use_grad(self, *args, **kwargs)
     31 try:
     32     torch.set_grad_enabled(self.defaults['differentiable'])
---> 33     ret = func(self, *args, **kwargs)
     34 finally:
     35     torch.set_grad_enabled(prev_grad)

File ~/anaconda3/envs/xgb3_ray2_pytorch/lib/python3.8/site-packages/torch/optim/adam.py:121, in Adam.step(self, closure)
    119 if closure is not None:
    120     with torch.enable_grad():
--> 121         loss = closure()
    123 for group in self.param_groups:
    124     params_with_grad = []

File ~/anaconda3/envs/xgb3_ray2_pytorch/lib/python3.8/site-packages/skorch/net.py:1077, in NeuralNet.train_step.<locals>.step_fn()
   1075 def step_fn():
   1076     self._zero_grad_optimizer()
-> 1077     step = self.train_step_single(batch, **fit_params)
   1078     step_accumulator.store_step(step)
   1080     self.notify(
   1081         'on_grad_computed',
   1082         named_parameters=TeeGenerator(self.get_all_learnable_params()),
   1083         batch=batch,
   1084         training=True,
   1085     )

File ~/anaconda3/envs/xgb3_ray2_pytorch/lib/python3.8/site-packages/skorch/net.py:977, in NeuralNet.train_step_single(self, batch, **fit_params)
    975 Xi, yi = unpack_data(batch)
    976 y_pred = self.infer(Xi, **fit_params)
--> 977 loss = self.get_loss(y_pred, yi, X=Xi, training=True)
    978 loss.backward()
    979 return {
    980     'loss': loss,
    981     'y_pred': y_pred,
    982 }

File ~/anaconda3/envs/xgb3_ray2_pytorch/lib/python3.8/site-packages/skorch/classifier.py:128, in NeuralNetClassifier.get_loss(self, y_pred, y_true, *args, **kwargs)
    126     eps = torch.finfo(y_pred.dtype).eps
    127     y_pred = torch.log(y_pred + eps)
--> 128 return super().get_loss(y_pred, y_true, *args, **kwargs)

File ~/anaconda3/envs/xgb3_ray2_pytorch/lib/python3.8/site-packages/skorch/net.py:1648, in NeuralNet.get_loss(self, y_pred, y_true, X, training)
   1619 """Return the loss for this batch.
   1620 
   1621 Parameters
   (...)
   1645 
   1646 """
   1647 y_true = to_tensor(y_true, device=self.device)
-> 1648 return self.criterion_(y_pred, y_true)

File ~/anaconda3/envs/xgb3_ray2_pytorch/lib/python3.8/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File ~/anaconda3/envs/xgb3_ray2_pytorch/lib/python3.8/site-packages/torch/nn/modules/loss.py:1174, in CrossEntropyLoss.forward(self, input, target)
   1173 def forward(self, input: Tensor, target: Tensor) -> Tensor:
-> 1174     return F.cross_entropy(input, target, weight=self.weight,
   1175                            ignore_index=self.ignore_index, reduction=self.reduction,
   1176                            label_smoothing=self.label_smoothing)

File ~/anaconda3/envs/xgb3_ray2_pytorch/lib/python3.8/site-packages/torch/nn/functional.py:3029, in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction, label_smoothing)
   3027 if size_average is not None or reduce is not None:
   3028     reduction = _Reduction.legacy_get_string(size_average, reduce)
-> 3029 return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)

RuntimeError: "nll_loss_forward_reduce_cuda_kernel_2d_index" not implemented for 'Float'

Could you please help how to fix it?

[BenjaminBossan]

Is it possible that the skorch code runs on GPU (what is the value of device?) and the pure PyTorch code runs on CPU? If so, please run the PyTorch code on GPU too and check if the same error occurs or not.