Raise error for 1D (size > 1) -> 0D parameter loads by dsashidh · Pull Request #166335 · pytorch/pytorch

dsashidh · 2025-10-27T20:22:16Z

Title

Fix load_state_dict: raise error for 1D (size > 1) -> 0D parameter loads

Summary

This PR fixes a bug where loading a 1D tensor (size > 1) into a scalar (0D) parameter would silently take the first element instead of raising an error. The fix preserves backward compatibility for 1D tensors of size 1 while catching genuine shape mismatches.

Motivation

Previously, loading a 1D tensor like torch.randn(32000) into a 0D scalar parameter would silently slice the first element, leading to silent data loss and potential bugs. This change ensures users get a clear error when there's a genuine shape mismatch.

Behavior change

Before:
1D tensor (any length) -> 0D scalar -> silently coerced using input_param[0]

After:

1D tensor (size == 1) -> 0D scalar -> allowed (backward compatibility)
1D tensor (size > 1) -> 0D scalar -> raises RuntimeError with size mismatch message

In torch/nn/modules/module.py, _load_from_state_dict, added input_param.shape[0] == 1 check to the backward compatibility condition to only allow single-element 1D tensors.

Tests

Added test_scalar_param_1d_tensor_raises to verify that loading 1D tensors of size > 1 raises an error, while size 1 loads successfully.

pytorch-bot · 2025-10-27T20:22:20Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/166335

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 8d2fa97 with merge base ed4aa44 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

dsashidh · 2025-10-27T20:22:24Z

@pytorchbot label "topic: not user facing"

morrison-turnansky · 2025-11-01T15:40:54Z

torch/nn/modules/module.py

@dsashidh It looks like this case was here for backwards compatibility, but from a long time ago.

If there is a decision to no longer support this backward compatibility, the cleaner fix would be to remove this if statement and let if fall back to if not is_param_lazy and input_param.shape != param.shape, where it will be caught in this case and currently has the same body.

Thanks for the feedback, I commented below.

mikaylagawarecki

I think my comment below might be an acceptable fix

mikaylagawarecki · 2025-11-03T18:36:01Z

torch/nn/modules/module.py

@@ -2436,7 +2436,11 @@ def _load_from_state_dict(
                    and len(param.shape) == 0
                    and len(input_param.shape) == 1


Perhaps add and input_param.shape[0] == 1, so the unexpected case will not fall into this if statement

Thanks for the feedback, I commented below.

dsashidh · 2025-11-03T22:01:41Z

Thank you for the feedback!
I see two possible approaches:
Option 1: Remove the entire backward compatibility block (@morrison-turnansky 's approach)
Option 2: Add and input_param.shape[0] == 1 to preserve backward compatibility for [1] -> scalar (@mikaylagawarecki's approach)

I wrote my test with the expectation that [1] -> scalar should raise an error (strict shape matching), which passes with approach 1 but not approach 2. Since this backward compatibility is from PyTorch 0.3 (2017), I'm inclined toward approach 1 for stricter shape matching.
@mikaylagawarecki : Is there a strong reason to preserve [1] -> scalar compatibility?

mikaylagawarecki · 2025-11-03T22:59:49Z

I think it should be fine to load a 1d tensor of size 1 into a scalar tensor, iiuc that's what adding the check I suggested would do, though do correct me if I'm wrong

morrison-turnansky · 2025-11-04T15:33:25Z

@dsashidh Go with what @mikaylagawarecki is suggesting.

dsashidh · 2025-11-04T16:08:50Z

I think it should be fine to load a 1d tensor of size 1 into a scalar tensor, iiuc that's what adding the check I suggested would do, though do correct me if I'm wrong

Hi @mikaylagawarecki thanks for clarifying! I've implemented your suggested check (and input_param.shape[0] == 1) and updated my test accordingly.
The fix now:
Allows [1] -> scalar (backward compatibility)
Raises error for [2+] -> scalar
Tests pass with this approach.

dsashidh · 2025-11-06T15:13:10Z

Hi @mikaylagawarecki I was seeing a MYPY lintrunner failure unrelated to my changes. I’ve rebased my branch on upstream/viable/strict which should hopefully align it with a clean CI baseline

mikaylagawarecki · 2025-11-07T01:52:40Z

@pytorchbot merge

pytorchmergebot · 2025-11-07T01:54:31Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

Summary: Fix in one OSS notebook where the state_dict was naively expanded. This [pytorch PR](pytorch/pytorch#166335) caused an error that was previously silently ignored. Differential Revision: D86884189

Summary: Pull Request resolved: #3079 Fix in one OSS notebook where the state_dict was naively expanded. This [pytorch PR](pytorch/pytorch#166335) caused an error that was previously silently ignored. Reviewed By: sdaulton Differential Revision: D86884189 fbshipit-source-id: 2c5ded01b17800a64da53b0722cfcc1ccac5e6eb

Fixes pytorch#165873 # Title Fix load_state_dict: raise error for 1D (size > 1) -> 0D parameter loads ## Summary This PR fixes a bug where loading a 1D tensor (size > 1) into a scalar (0D) parameter would silently take the first element instead of raising an error. The fix preserves backward compatibility for 1D tensors of size 1 while catching genuine shape mismatches. ## Motivation Previously, loading a 1D tensor like torch.randn(32000) into a 0D scalar parameter would silently slice the first element, leading to silent data loss and potential bugs. This change ensures users get a clear error when there's a genuine shape mismatch. ## Behavior change Before: 1D tensor (any length) -> 0D scalar -> silently coerced using input_param[0] After: - 1D tensor (size == 1) -> 0D scalar -> allowed (backward compatibility) - 1D tensor (size > 1) -> 0D scalar -> raises RuntimeError with size mismatch message In torch/nn/modules/module.py, _load_from_state_dict, added input_param.shape[0] == 1 check to the backward compatibility condition to only allow single-element 1D tensors. ## Tests Added test_scalar_param_1d_tensor_raises to verify that loading 1D tensors of size > 1 raises an error, while size 1 loads successfully. Pull Request resolved: pytorch#166335 Approved by: https://github.com/mikaylagawarecki

dsashidh requested review from albanD, jbschlosser and mikaylagawarecki as code owners October 27, 2025 20:22

pytorch-bot bot added the topic: not user facing topic category label Oct 27, 2025

pytorchbot added the open source label Oct 27, 2025

soulitzer added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Oct 28, 2025

morrison-turnansky reviewed Nov 1, 2025

View reviewed changes

mikaylagawarecki reviewed Nov 3, 2025

View reviewed changes

albanD removed their request for review November 3, 2025 22:05

dsashidh changed the title ~~Raise error on 1D - > 0D parameter loads in load_state_dict~~ Raise error for 1D (size > 1) -> 0D parameter loads Nov 4, 2025

dsashidh added 2 commits November 6, 2025 15:08

fixes scalar-1D tensor mismatch error

9965764

fix load_state_dict shape mismatch

8d2fa97

dsashidh force-pushed the fix_load_state_dict_shape_mismatch branch from 6b097c0 to 8d2fa97 Compare November 6, 2025 15:09

mikaylagawarecki approved these changes Nov 6, 2025

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Nov 7, 2025

mikaylagawarecki added release notes: nn release notes category topic: bug fixes topic category and removed topic: not user facing topic category labels Nov 7, 2025

pytorchmergebot added the merging label Nov 7, 2025

pytorchmergebot added the Merged label Nov 7, 2025

pytorchmergebot closed this in 8a72188 Nov 7, 2025

pytorchmergebot removed the merging label Nov 7, 2025

Balandat mentioned this pull request Nov 11, 2025

prevent warning meta-pytorch/botorch#3078

Closed

hvarfner mentioned this pull request Nov 12, 2025

Fix of failing OSS NB test meta-pytorch/botorch#3079

Closed

		@@ -2436,7 +2436,11 @@ def _load_from_state_dict(
		and len(param.shape) == 0
		and len(input_param.shape) == 1

Conversation

dsashidh commented Oct 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Title

Summary

Motivation

Behavior change

Tests

Uh oh!

pytorch-bot bot commented Oct 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/166335

✅ No Failures

Uh oh!

dsashidh commented Oct 27, 2025

Uh oh!

morrison-turnansky Nov 1, 2025

Choose a reason for hiding this comment

Uh oh!

dsashidh Nov 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mikaylagawarecki left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mikaylagawarecki Nov 3, 2025

Choose a reason for hiding this comment

Uh oh!

dsashidh Nov 3, 2025

Choose a reason for hiding this comment

Uh oh!

dsashidh commented Nov 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mikaylagawarecki commented Nov 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

morrison-turnansky commented Nov 4, 2025

Uh oh!

dsashidh commented Nov 4, 2025

Uh oh!

dsashidh commented Nov 6, 2025

Uh oh!

mikaylagawarecki commented Nov 7, 2025

Uh oh!

pytorchmergebot commented Nov 7, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

dsashidh commented Oct 27, 2025 •

edited

Loading

pytorch-bot bot commented Oct 27, 2025 •

edited

Loading

dsashidh Nov 3, 2025 •

edited

Loading

mikaylagawarecki left a comment •

edited

Loading

dsashidh commented Nov 3, 2025 •

edited

Loading

mikaylagawarecki commented Nov 3, 2025 •

edited

Loading