Training fails with tensor size mismatch (5 vs 10 channels) after migrating project from 1.4.x to 1.5

[hlanovoi]

Bug description

Training fails with RuntimeError: The size of tensor a (5) must match the size of tensor b (10) on a project migrated from SLEAP 1.4.x to 1.5. The model is correctly configured for 5 channels (matching the 5-node skeleton), but the training data loader reports 10 channels, causing a tensor broadcast error.

Expected behaviour

Training should succeed with a 5-channel model matching the 5-node skeleton in the project.

Actual behaviour

Training fails immediately with:

RuntimeError: The size of tensor a (5) must match the size of tensor b (10) at non-singleton dimension 1

The error occurs in sleap_nn/training/lightning_modules.py, line 537 during loss computation. The model head correctly outputs 5 channels, but the target tensor from the data pipeline has 10 channels.

Your personal set up

• OS: macOS (Apple Silicon - M-series chip)
• Version(s): SLEAP 1.5.x, Python 3.13
• SLEAP installation method: uv tool install "sleap[nn]"

# paste relevant logs here, if any

/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/torch/nn/modules/loss.py:616: UserWarning: Using a target size (torch.Size([4, 10, 240, 320])) that is different to the input size (torch.Size([4, 5, 240, 320])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.
return F.mse_loss(input, target, reduction=self.reduction)

Traceback (most recent call last):
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/sleap_nn/training/lightning_modules.py", line 537, in training_step
loss = nn.MSELoss()(y_preds, y)
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/torch/nn/modules/loss.py", line 616, in forward
return F.mse_loss(input, target, reduction=self.reduction)
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/torch/nn/functional.py", line 3868, in mse_loss
expanded_input, expanded_target = torch.broadcast_tensors(input, target)
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/torch/functional.py", line 77, in broadcast_tensors
return _VF.broadcast_tensors(tensors)
RuntimeError: The size of tensor a (5) must match the size of tensor b (10) at non-singleton dimension 1

**Training output shows:**
- Model head: `Conv2d(32, 5, kernel_size=(1, 1))` 
- Target tensor: `torch.Size([4, 10, 240, 320])` 

## Screenshots
N/A

## How to reproduce
1. Open a SLEAP project that was created in v1.4.x (now opened in v1.5)
2. Project has a single skeleton with 5 nodes
3. Attempt to train a single instance model using the SLEAP GUI (Predict → Run Training → Single Instance)
4. Training fails immediately with tensor size mismatch error

**Additional diagnostic steps I've run through:**
- Verified only 1 skeleton with 5 nodes exists using `sleap_io.load_slp()` 
- Exported clean package with only hand-labeled frames (Predict → Export Labels Package → "Labeled frames")
- Created new project from clean package and attempted training
- Error persists even with fresh export

[gitttt-1234]

Hi @hlanovoi !

Could you please share the training_config.yaml (this would be saved in the ckpt folder created for the training run)? Could you also paste the entire log from the terminal?

Thanks,

Divya

[hlanovoi]

I appreciate the prompt response @gitttt-1234 !

Here is the training_config.yaml from the failed training:

training_config.yaml

Here is what I see in my terminal:

Polling: /Users/HL801/Desktop/SLEAP_anaylsis/MAIN_Projects/4_cup_apparatus(2)/models/251214_183623.single_instance.n=4700/viz/validation.*.png
Start training single_instance...
['sleap-nn-train', '--config-name', '251214_183623_config', '--config-dir', '/var/folders/k4/szpsbld96q97g__7r0k2xn7r0000gp/T/tmpgsdyfx4o']
Input config:

data_config:
train_labels_path:

• /Users/HL801/Desktop/SLEAP_anaylsis/MAIN_Projects/4_cup_apparatus(2)/4cup2.v004_w2cup.v013.slp
  val_labels_path: null
  validation_fraction: 0.1
  test_file_path: null
  provider: LabelsReader
  user_instances_only: true
  data_pipeline_fw: torch_dataset
  cache_img_path: null
  use_existing_imgs: false
  delete_cache_imgs_after_training: true
  preprocessing:
  ensure_rgb: false
  ensure_grayscale: false
  max_height: 480
  max_width: 640
  scale: 1.0
  crop_size: null
  min_crop_size: 100
  use_augmentations_train: true
  augmentation_config:
  intensity:
  uniform_noise_min: 0.0
  uniform_noise_max: 1.0
  uniform_noise_p: 0.0
  gaussian_noise_mean: 5.0
  gaussian_noise_std: 1.0
  gaussian_noise_p: 0.0
  contrast_min: 0.5
  contrast_max: 2.0
  contrast_p: 0.0
  brightness_min: 0.0
  brightness_max: 2.0
  brightness_p: 0.0
  geometric:
  rotation_min: -180.0
  rotation_max: 180.0
  scale_min: 0.9
  scale_max: 1.1
  translate_width: 0.0
  translate_height: 0.0
  affine_p: 1.0
  erase_scale_min: 0.0001
  erase_scale_max: 0.01
  erase_ratio_min: 1.0
  erase_ratio_max: 1.0
  erase_p: 0.0
  mixup_lambda_min: 0.01
  mixup_lambda_max: 0.05
  mixup_p: 0.0
  skeletons: []
  model_config:
  init_weights: default
  pretrained_backbone_weights: null
  pretrained_head_weights: null
  backbone_config:
  unet:
  in_channels: 1
  kernel_size: 3
  filters: 16
  filters_rate: 2.0
  max_stride: 16
  stem_stride: null
  middle_block: true
  up_interpolate: true
  stacks: 1
  convs_per_block: 2
  output_stride: 2
  convnext: null
  swint: null
  head_configs:
  single_instance:
  confmaps:
  part_names:
  - Nose
  - Left ear
  - Right ear
  - Center point
  - Tail point
  sigma: 2.5
  output_stride: 2
  centroid: null
  centered_instance: null
  bottomup: null
  multi_class_bottomup: null
  multi_class_topdown: null
  total_params: null
  trainer_config:
  train_data_loader:
  batch_size: 4
  num_workers: 0
  shuffle: true
  val_data_loader:
  batch_size: 4
  num_workers: 0
  shuffle: false
  model_ckpt:
  save_top_k: 1
  save_last: false
  trainer_devices: null
  trainer_device_indices: null
  trainer_accelerator: auto
  profiler: null
  trainer_strategy: auto
  enable_progress_bar: true
  min_train_steps_per_epoch: 200
  train_steps_per_epoch: 602
  visualize_preds_during_training: true
  keep_viz: false
  max_epochs: 200
  seed: null
  use_wandb: false
  save_ckpt: true
  ckpt_dir: /Users/HL801/Desktop/SLEAP_anaylsis/MAIN_Projects/4_cup_apparatus(2)/models
  run_name: 251214_183623.single_instance.n=4700
  resume_ckpt_path: null
  wandb:
  entity: ''
  project: ''
  name: ''
  save_viz_imgs_wandb: false
  api_key: ''
  wandb_mode: null
  prv_runid: ''
  group: ''
  current_run_id: null
  optimizer_name: Adam
  optimizer:
  lr: 0.0001
  amsgrad: false
  lr_scheduler:
  step_lr: null
  reduce_lr_on_plateau:
  threshold: 1.0e-06
  threshold_mode: abs
  cooldown: 3
  patience: 5
  factor: 0.5
  min_lr: 1.0e-08
  early_stopping:
  min_delta: 1.0e-08
  patience: 10
  stop_training_on_plateau: true
  online_hard_keypoint_mining:
  online_mining: false
  hard_to_easy_ratio: 2.0
  min_hard_keypoints: 2
  max_hard_keypoints: null
  loss_scale: 5.0
  zmq:
  controller_port: 9000
  controller_polling_timeout: 10
  publish_port: 9001
  name: ''
  description: ''
  sleap_nn_version: 0.0.5
  filename: ''

2025-12-14 18:36:26 | INFO | sleap_nn.train:run_training:24 | Started training at: 2025-12-14 18:36:26.597498
2025-12-14 18:36:26 | INFO | sleap_nn.training.model_trainer:_setup_train_val_labels:216 | Creating train-val split...
2025-12-14 18:36:27 | INFO | sleap_nn.training.model_trainer:_setup_train_val_labels:261 | # Train Labeled frames: 4230
2025-12-14 18:36:27 | INFO | sleap_nn.training.model_trainer:_setup_train_val_labels:262 | # Val Labeled frames: 470
2025-12-14 18:36:27 | INFO | sleap_nn.training.model_trainer:setup_config:510 | Setting up config...
2025-12-14 18:36:27 | INFO | sleap_nn.training.model_trainer:_verify_model_input_channels:416 | Updating backbone in_channels to 3 based on the input image channels.
2025-12-14 18:36:27 | INFO | sleap_nn.training.model_trainer:train:849 | Setting up for training...
2025-12-14 18:36:27 | INFO | sleap_nn.training.model_trainer:_setup_model_ckpt_dir:573 | Setting up model ckpt dir: /Users/HL801/Desktop/SLEAP_anaylsis/MAIN_Projects/4_cup_apparatus(2)/models/251214_183623.single_instance.n=4700...
2025-12-14 18:36:27 | INFO | sleap_nn.training.model_trainer:train:864 | Setting up visualization train and val datasets...
2025-12-14 18:36:27 | INFO | sleap_nn.training.model_trainer:train:868 | Setting up Trainer...
2025-12-14 18:36:27 | INFO | sleap_nn.training.model_trainer:_setup_loggers_callbacks:645 | Setting up callbacks and loggers...
2025-12-14 18:36:27 | INFO | sleap_nn.training.callbacks:init:228 | Training controller subscribed to: tcp://127.0.0.1:9000 (topic: )
2025-12-14 18:36:27 | INFO | sleap_nn.training.callbacks:init:287 | ProgressReporterZMQ publishing to tcp://127.0.0.1:9001 for ''
2025-12-14 18:36:27 | INFO | sleap_nn.training.model_trainer:train:897 | Trainer devices: auto
INFO: GPU available: True (mps), used: True
INFO:lightning.pytorch.utilities.rank_zero:GPU available: True (mps), used: True
INFO: TPU available: False, using: 0 TPU cores
INFO:lightning.pytorch.utilities.rank_zero:TPU available: False, using: 0 TPU cores
2025-12-14 18:36:27 | INFO | sleap_nn.training.model_trainer:train:950 | Training on 1 device(s)
2025-12-14 18:36:27 | INFO | sleap_nn.training.model_trainer:train:951 | Training on mps:0 accelerator
2025-12-14 18:36:27 | INFO | sleap_nn.training.model_trainer:train:955 | Setting up lightning module for single_instance model...
2025-12-14 18:36:27 | INFO | sleap_nn.training.model_trainer:train:959 | Backbone model: UNet(
(encoders): ModuleList(
(0): Encoder(
(encoder_stack): ModuleList(
(0): SimpleConvBlock(
(blocks): Sequential(
(stack0_enc0_conv0): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=same)
(stack0_enc0_act0_relu): ReLU()
(stack0_enc0_conv1): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=same)
(stack0_enc0_act1_relu): ReLU()
)
)
(1): SimpleConvBlock(
(blocks): Sequential(
(stack0_enc1_pool): MaxPool2dWithSamePadding(kernel_size=2, stride=2, padding=same, dilation=1, ceil_mode=False)
(stack0_enc1_conv0): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=same)
(stack0_enc1_act0_relu): ReLU()
(stack0_enc1_conv1): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=same)
(stack0_enc1_act1_relu): ReLU()
)
)
(2): SimpleConvBlock(
(blocks): Sequential(
(stack0_enc2_pool): MaxPool2dWithSamePadding(kernel_size=2, stride=2, padding=same, dilation=1, ceil_mode=False)
(stack0_enc2_conv0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=same)
(stack0_enc2_act0_relu): ReLU()
(stack0_enc2_conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=same)
(stack0_enc2_act1_relu): ReLU()
)
)
(3): SimpleConvBlock(
(blocks): Sequential(
(stack0_enc3_pool): MaxPool2dWithSamePadding(kernel_size=2, stride=2, padding=same, dilation=1, ceil_mode=False)
(stack0_enc3_conv0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=same)
(stack0_enc3_act0_relu): ReLU()
(stack0_enc3_conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=same)
(stack0_enc3_act1_relu): ReLU()
)
)
(4): Sequential(
(stack0_enc4_last_pool): MaxPool2dWithSamePadding(kernel_size=2, stride=2, padding=same, dilation=1, ceil_mode=False)
)
)
)
)
(decoders): ModuleList(
(0): Decoder(
(decoder_stack): ModuleList(
(0): SimpleUpsamplingBlock(
(blocks): Sequential(
(stack0_dec0_s16_to_s8_interp_bilinear): Upsample(scale_factor=2.0, mode='bilinear')
(stack0_dec0_s16_to_s8_refine_conv0): Conv2d(384, 128, kernel_size=(3, 3), stride=(1, 1), padding=same)
(stack0_dec0_s16_to_s8_refine_conv0_act_relu): ReLU()
(stack0_dec0_s16_to_s8_refine_conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=same)
(stack0_dec0_s16_to_s8_refine_conv1_act_relu): ReLU()
)
)
(1): SimpleUpsamplingBlock(
(blocks): Sequential(
(stack0_dec1_s8_to_s4_interp_bilinear): Upsample(scale_factor=2.0, mode='bilinear')
(stack0_dec1_s8_to_s4_refine_conv0): Conv2d(192, 64, kernel_size=(3, 3), stride=(1, 1), padding=same)
(stack0_dec1_s8_to_s4_refine_conv0_act_relu): ReLU()
(stack0_dec1_s8_to_s4_refine_conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=same)
(stack0_dec1_s8_to_s4_refine_conv1_act_relu): ReLU()
)
)
(2): SimpleUpsamplingBlock(
(blocks): Sequential(
(stack0_dec2_s4_to_s2_interp_bilinear): Upsample(scale_factor=2.0, mode='bilinear')
(stack0_dec2_s4_to_s2_refine_conv0): Conv2d(96, 32, kernel_size=(3, 3), stride=(1, 1), padding=same)
(stack0_dec2_s4_to_s2_refine_conv0_act_relu): ReLU()
(stack0_dec2_s4_to_s2_refine_conv1): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=same)
(stack0_dec2_s4_to_s2_refine_conv1_act_relu): ReLU()
)
)
)
)
)
(middle_blocks): ModuleList(
(0): SimpleConvBlock(
(blocks): Sequential(
(stack0_enc5_middle_expand_conv0): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=same)
(stack0_enc5_middle_expand_act0_relu): ReLU()
)
)
(1): SimpleConvBlock(
(blocks): Sequential(
(stack0_enc6_middle_contract_conv0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=same)
(stack0_enc6_middle_contract_act0_relu): ReLU()
)
)
)
)
2025-12-14 18:36:27 | INFO | sleap_nn.training.model_trainer:train:960 | Head model: ModuleList(
(0): Sequential(
(SingleInstanceConfmapsHead): Sequential(
(0): Conv2d(32, 5, kernel_size=(1, 1), stride=(1, 1), padding=same)
(1): Identity()
)
)
)
2025-12-14 18:36:27 | INFO | sleap_nn.training.model_trainer:train:962 | Total model parameters: 1953525
2025-12-14 18:36:28 | INFO | sleap_nn.training.model_trainer:train:967 | Input image shape: torch.Size([1, 3, 480, 640])
/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/torch/utils/data/dataloader.py:684: UserWarning: 'pin_memory' argument is set as true but not supported on MPS now, then device pinned memory won't be used.
warnings.warn(warn_msg)
2025-12-14 18:36:28 | INFO | sleap_nn.training.model_trainer:train:1021 | Finished trainer set up. [0.7s]
2025-12-14 18:36:28 | INFO | sleap_nn.training.model_trainer:train:1024 | Starting training loop...
/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/callbacks/model_checkpoint.py:881: Checkpoint directory /Users/HL801/Desktop/SLEAP_anaylsis/MAIN_Projects/4_cup_apparatus(2)/models/251214_183623.single_instance.n=4700 exists and is not empty.
┏━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━┳━━━━━━━┓
┃ ┃ Name ┃ Type ┃ Params ┃ Mode ┃ FLOPs ┃
┡━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━╇━━━━━━━┩
│ 0 │ model │ Model │ 2.0 M │ train │ 0 │
│ 1 │ single_instance_inf_layer │ SingleInstanceInferenceModel │ 0 │ train │ 0 │
└───┴───────────────────────────┴──────────────────────────────┴────────┴───────┴───────┘
Trainable params: 2.0 M
Non-trainable params: 0
Total params: 2.0 M
Total estimated model params size (MB): 7
Modules in train mode: 73
Modules in eval mode: 0
Total FLOPs: 0
/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/trainer/connectors/data_connector.p
y:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the
num_workers argumenttonum_workers=11in theDataLoaderto improve performance. /Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/core/module.py:522: You called self.log('learning_rate', ..., logger=True)but have no logger configured. You can enable one by doing Trainer(logger=ALogger(...))/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/core/module.py:522: You called self.log('val_loss', ..., logger=True)but have no logger configured. You can enable one by doing Trainer(logger=ALogger(...))/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/core/module.py:522: You called self.log('val_time', ..., logger=True)but have no logger configured. You can enable one by doing Trainer(logger=ALogger(...))/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/trainer/connectors/data_connector.p y:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the num_workers argument to num_workers=11 in the DataLoader to improve performance.
/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/torch/nn/modules/loss.py:616: UserWarning: Using a
target size (torch.Size([4, 10, 240, 320])) that is different to the input size (torch.Size([4, 5, 240, 320])). This will
likely lead to incorrect results due to broadcasting. Please ensure they have the same size.
return F.mse_loss(input, target, reduction=self.reduction)
Epoch 0/199 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0/602 0:00:00 • -:--:-- 0.00it/s learning_rate_step: 0.000
val_loss_step: 0.010
learning_rate_epoch: 0.000
val_loss_epoch: 0.010
2025-12-14 18:36:28 | INFO | sleap_nn.training.model_trainer:train:1037 | Finished training loop. [0.0 min]
2025-12-14 18:36:28 | INFO | sleap_nn.training.model_trainer:train:1064 | Deleting viz folder at /Users/HL801/Desktop/SLEAP_anaylsis/MAIN_Projects/4_cup_apparatus(2)/models/251214_183623.single_instance.n=4700/viz...
Traceback (most recent call last):
File "/Users/HL801/.local/bin/sleap-nn-train", line 10, in
sys.exit(train())
~~~~~^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/click/core.py", line 1485, in call
return self.main(*args, **kwargs)
~~~~~~~~~^^^^^^^^^^^^^^^^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/click/core.py", line 1406, in main
rv = self.invoke(ctx)
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/click/core.py", line 1269, in invoke
return ctx.invoke(self.callback, **ctx.params)
~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/click/core.py", line 824, in invoke
return callback(*args, **kwargs)
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/sleap/cli.py", line 87, in train
run_training(cfg)
~~~~~~~~~~~~^^^^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/sleap_nn/train.py", line 27, in run_training
trainer.train()
~~~~~~~~~~~~~^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/sleap_nn/training/model_trainer.py", line 1026, in train
self.trainer.fit(
~~~~~~~~~~~~~~~~^
self.lightning_model,
^^^^^^^^^^^^^^^^^^^^^
...<2 lines>...
ckpt_path=self.config.trainer_config.resume_ckpt_path,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
)
^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/trainer/trainer.py", line 584, in fit
call._call_and_handle_interrupt(
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
self,
^^^^^
...<6 lines>...
weights_only,
^^^^^^^^^^^^^
)
^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/trainer/call.py", line 49, in _call_and_handle_interrupt
return trainer_fn(*args, **kwargs)
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/trainer/trainer.py", line 630, in _fit_impl
self._run(model, ckpt_path=ckpt_path, weights_only=weights_only)
~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/trainer/trainer.py", line 1079, in _run
results = self._run_stage()
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/trainer/trainer.py", line 1123, in _run_stage
self.fit_loop.run()
~~~~~~~~~~~~~~~~~^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/loops/fit_loop.py", line 217, in run
self.advance()
~~~~~~~~~~~~^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/loops/fit_loop.py", line 465, in advance
self.epoch_loop.run(self._data_fetcher)
~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 153, in run
self.advance(data_fetcher)
~~~~~~~~~~~~^^^^^^^^^^^^^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 352, in advance
batch_output = self.automatic_optimization.run(trainer.optimizers[0], batch_idx, kwargs)
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 192, in run
self._optimizer_step(batch_idx, closure)
~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 270, in _optimizer_step
call._call_lightning_module_hook(
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
trainer,
^^^^^^^^
...<4 lines>...
train_step_and_backward_closure,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
)
^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/trainer/call.py", line 177, in _call_lightning_module_hook
output = fn(*args, **kwargs)
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/core/module.py", line 1368, in optimizer_step
optimizer.step(closure=optimizer_closure)
~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/core/optimizer.py", line 154, in step
step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/strategies/strategy.py", line 239, in optimizer_step
return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/plugins/precision/precision.py", line 123, in optimizer_step
return optimizer.step(closure=closure, **kwargs)
~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/torch/optim/optimizer.py", line 516, in wrapper
out = func(*args, **kwargs)
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/torch/optim/optimizer.py", line 81, in _use_grad
ret = func(*args, **kwargs)
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/torch/optim/adam.py", line 226, in step
loss = closure()
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/plugins/precision/precision.py", line 109, in _wrap_closure
closure_result = closure()
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 146, in call
self._result = self.closure(*args, **kwargs)
~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
return func(*args, **kwargs)
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 131, in closure
step_output = self._step_fn()
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 319, in _training_step
training_step_output = call._call_strategy_hook(trainer, "training_step", *kwargs.values())
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/trainer/call.py", line 329, in _call_strategy_hook
output = fn(*args, **kwargs)
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/lightning/pytorch/strategies/strategy.py", line 391, in training_step
return self.lightning_module.training_step(*args, **kwargs)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/sleap_nn/training/lightning_modules.py", line 537, in training_step
loss = nn.MSELoss()(y_preds, y)
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
return forward_call(*args, **kwargs)
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/torch/nn/modules/loss.py", line 616, in forward
return F.mse_loss(input, target, reduction=self.reduction)
~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/torch/nn/functional.py", line 3868, in mse_loss
expanded_input, expanded_target = torch.broadcast_tensors(input, target)
~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^
File "/Users/HL801/.local/share/uv/tools/sleap/lib/python3.13/site-packages/torch/functional.py", line 77, in broadcast_tensors
return _VF.broadcast_tensors(tensors) # type: ignore[attr-defined]
~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^
RuntimeError: The size of tensor a (5) must match the size of tensor b (10) at non-singleton dimension 1
Run Path: /Users/HL801/Desktop/SLEAP_anaylsis/MAIN_Projects/4_cup_apparatus(2)/models/251214_183623.single_instance.n=4700

Thanks,
Hunter

[gitttt-1234]

Hi @hlanovoi !

Thank you for sharing the logs, the model seems right, could you please share your slp file here to help us debug better?

Thanks,

Divya