Skip to content

Relax Hessian dataloader batch unpacking for common formats - #17

Open
kabishou11 wants to merge 1 commit into
THUDM:mainfrom
kabishou11:fix/hessian-batch-unpack
Open

kabishou11 wants to merge 1 commit into
THUDM:mainfrom
kabishou11:fix/hessian-batch-unpack

Conversation

@kabishou11

Copy link
Copy Markdown

Summary

Fixes #10.

hessian.dataloader_hv_product hard-coded for _, inputs, targets in self.data, so the common PyTorch (inputs, targets) DataLoader convention raised ValueError: too many values to unpack.

Add unpack_batch that accepts:

  • (inputs, targets) pairs
  • legacy (index, inputs, targets) triplets
  • simple dict batches with inputs/targets (or input/target)

Test plan

  • uv run --with pytest pytest -q tests/test_hessian_batch_unpack.py — 4 passed

dataloader_hv_product previously required a 3-tuple
(index, inputs, targets), so standard PyTorch (inputs, targets)
DataLoaders raised ValueError. Add unpack_batch that accepts 2-tuples,
legacy 3-tuples, and simple dict batches, and cover both paths with
unit tests.

Fixes THUDM#10

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Hard-coded batch-triplet unpack breaks custom DataLoader formats

1 participant