Update deeplabcut/pose_estimation_pytorch/runners/train.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: xiu-cs

deeplabcut/pose_estimation_pytorch/runners/train.py

@@ -439,7 +439,7 @@ def step(
         if mode == "train":
             self.optimizer.zero_grad()
 
-        inputs = batch["image"] # [B,3, 256, 256]
+        inputs = batch["image"] # [batch_size, channels, height, width]
         inputs = inputs.to(self.device).float()
         if 'cond_keypoints' in batch['context']:
             cond_kpts = batch['context']['cond_keypoints']