Parte 03 - Aula - 02 - RuntimeError: "nll_loss_forward_reduce_cuda_kernel_2d_index" not implemented for 'Float'

Quando executo o código

for i in range(100):
  # Forward
  pred = net(X)
  loss = criterion(pred,Y)

  #Backward
  loss.backward()
  optimizer.step()
  if i % 10 == 0:
    plt.figure()
    plot_boundary(data, targets, net)

o colab emite o seguinte erro:

RuntimeError Traceback (most recent call last) in <cell line: 1>() 2 # Forward 3 pred = net(X) ----> 4 loss = criterion(pred,Y) 5 6 #Backward

3 frames /usr/local/lib/python3.10/dist-packages/torch/nn/functional.py in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction, label_smoothing) 3051 if size_average is not None or reduce is not None: 3052 reduction = _Reduction.legacy_get_string(size_average, reduce) -> 3053 return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing) 3054 3055

RuntimeError: "nll_loss_forward_reduce_cuda_kernel_2d_index" not implemented for 'Float'

alguém já passou por este problema?