UnboundLocalError: local variable referenced before assignment from import torch.optim._multi_tensor

Hi all,

I am facing the error

UnboundLocalError: local variable 'torch' referenced before assignment

from running the following snippet

import torch
import torch.nn as nn

class ASmallModel(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

        print(f"from small model: {torch}")

        import torch.optim._multi_tensor


def main():
    mym = ASmallModel()


if __name__ == "__main__":
    main()

The solution is that moving the import torch.optim._multi_tensor to the start of the script. However, I am just curious how it causes the error when I do the import after print(f"from small model: {torch}")

Thanks a lot!

Best,
Jiawei

My torch version is

torch                    2.0.1
torchvision              0.15.2