Update train.py

fffffgggg54 · Aug 7, 2024 · 969955e · 969955e
1 parent d77a8ad
commit 969955e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.py b/train.py
@@ -1006,7 +1006,7 @@ def _forward():
  with amp_autocast():
  output = model(input)
  loss = loss_fn(output, target)
- acc1 = utils.accuracy(output, target.topk(1, 1), topk=(1,))[0]
+ acc1 = utils.accuracy(output, target.topk(1, 1)[1], topk=(1,))[0]
 
  if accum_steps > 1:
  loss /= accum_steps