Fix state_dict loading for models trained with torch.compile

[inorganicwriter]

Problem

When a model is trained with compile: True (the current default in
config.py), torch.compile prepends the _orig_mod. prefix to every
parameter key in the saved state_dict. Loading those weights into a
non-compiled model fails with a key-mismatch error. This commonly
happens when:

• Running run evaluate / run infer on a host where torch.compile
  is disabled (e.g. Windows without MSVC cl.exe).
• Setting compile: False in the config for inference while the
  checkpoint was produced with compile: True.
• Sharing checkpoints across compiled/non-compiled training runs.

Fix

In BaseTester._load_model_weights, strip the _orig_mod. prefix from
all keys before calling load_state_dict when its presence is detected.
The check is opt-in (only triggered when the prefix is actually there),
so loading non-compiled weights is unchanged.

Testing

• Loaded a compile=True checkpoint with compile=False set in the
  eval config: previously raised RuntimeError, now loads cleanly.
• Loaded a compile=False checkpoint: behavior unchanged.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[inorganicwriter]

@google-cla-bot check