Problemas para converter um tensor do torch para numpy array

Não consigo executar o foward da aula 2, sempre que eu chamo a função ele retorna o erro:

TypeError: can't convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.

eu reescrevi e testei a função linha por linha o erro acontece na linha loss_epoca = np.array(loss_epoca)

def forward(X, Y, etapa, epoca):

  acuracia = 0
  loss_epoca = []
  for dado, rotulo in zip(X, Y):
    dado   = dado.to(device)
    rotulo = rotulo.to(device)

    saida = modelo(dado) # B x C = 1 x 18
    loss = criterio(saida, rotulo)
    loss_epoca.append(loss)

    _, pred = torch.max(saida, axis=-1)
    acuracia += 1 if pred[0].item() == rotulo[0].item() else 0

    if etapa == 'treino':
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

    loss_epoca = np.array(loss_epoca) # <<<- O erro acontece aqui, 
    acuracia   = acuracia/float(len(X))

    print('\n', '*' * 15 + etapa + '*' * 15)
    print('Epoca: {:}, Loss: {:.2f}, Acc: {:.2f}'.format(epoca, 
                                                         loss.mean(), 
                                                         loss.std(), 
                                                         acuracia*100 ))
    return loss.mean(), acuracia

eu tentei usar algumas funções pra resolver tipo .cpu() e .numpy() mas nenhum soluciona o erro .