Erro na execução do Forward

Olá, Estou com problema no código, eu estou na aula 3 parte 2, estou fazendo tudo como no vídeo, passo a passo, mas na parte de Forward deu erro, eu verifiquei se ta certo e esta como no vídeo, cliquei no link e o colab da pra ver resultados do stack overflow e da apenas 3 resultados q não condizem com o problema, re vi o vídeo e o anterior para conferir os códigos, até ai ta tudo ok, já finalizei a cessão do colab e abri um nova em uma nova aba e nada, mesmo erro.

# Forward
pred = net(X)
loss = criterion(pred, Y)

# Backward
loss.backward()
optimizer.step()

plot_boundary(data, target, net)

ValueError                                Traceback (most recent call last)
<ipython-input-77-56c9cb3de796> in <module>()
      1 # Forward
      2 pred = net(X)
----> 3 loss = criterion(pred, Y)
      4 
      5 # Backward

3 frames
/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py in nll_loss(input, target, weight, size_average, ignore_index, reduce, reduction)
   2214     if input.size(0) != target.size(0):
   2215         raise ValueError('Expected input batch_size ({}) to match target batch_size ({}).'
-> 2216                          .format(input.size(0), target.size(0)))
   2217     if dim == 2:
   2218         ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)

ValueError: Expected input batch_size (178) to match target batch_size (442).