Hi Jake!



 Jake_Eum:

    criterion = nn.CrossEntropyLoss()
    loss = criterion(output[0],target[0])



I see two issues:
First, pytorch (models, losses, etc.) works with batches of samples.
criterion (output[0], target[0]) looks like you are passing
a single sample from a batch to criterion.  This will break things.
Second, the shapes of output and target don’t match properly for
use with CrossEntropyLoss.  CrossEntropyLoss expects your
output to have shape [nBatch, nClass] and target to have
shape [nBatch] (no class dimension).  target should consist of
integer class labels that run from 0 to nClass - 1.
One final comment:  Your target looks like it is only 0s and 1s, so
it looks like you are working with a binary classification problem.
You can treat a binary problem as a two-class multi-class problem.
Doing so is fine, and makes perfect sense, but you might prefer
treating this explicitly as a binary problem and use something like
BCEWithLogitsLoss as your loss criterion.
Best.
K. Frank