Change weights return type to Mapping

[datumbox]

Cherrypick of #6097

Either both go in, or none goes in

[datumbox]

Closing this. The issue is that the wrong version of PyTorch core is being installed:

Looking in indexes: https://pypi.org/simple, https://download.pytorch.org/whl/test/cpu/torch_test.html
Collecting torch
  Downloading torch-1.11.0-cp37-cp37m-manylinux1_x86_64.whl (750.6 MB)
Requirement already satisfied: typing-extensions in /home/circleci/.pyenv/versions/3.7.13/lib/python3.7/site-packages (from torch) (4.1.1)
Installing collected packages: torch
Successfully installed torch-1.11.0
WARNING: You are using pip version 22.0.4; however, version 22.1.1 is available.
You should consider upgrading via the '/home/circleci/.pyenv/versions/3.7.13/bin/python3.7 -m pip install --upgrade pip' command.

We will open a ticket to inform the RelEng team.