Fix of failing OSS NB test (meta-pytorch#3079)

Summary:
Pull Request resolved: meta-pytorch#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

Author: Carl Hvarfner