torch.export: cannot instantiate Dim from REPL (#111231)
Summary:
```
In [1]: import torch
...: torch.export.Dim('foo', min=1, max=16)
---------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
Cell In[1], line 2
1 import torch
----> 2 torch.export.Dim('foo', min=1, max=16)
File /..../torch/export/__init__.py:319, in Dim(name, min, max)
317 assert _max > _min, f"Cannot create Dim with inconsistent min={min}, max={max}"
318 dim = _Dim(name, (int,), {"min": _min, "max": _max})
--> 319 dim.__module__ = inspect.getmodule(inspect.stack()[1][0]).__name__ # type: ignore[union-attr]
320 return dim
AttributeError: 'NoneType' object has no attribute '__name__'
```
Test Plan: Repeat above repro
Differential Revision: D50275165
Pull Request resolved: https://github.com/pytorch/pytorch/pull/111231
Approved by: https://github.com/avikchaudhuri, https://github.com/angelayi
diff --git a/torch/export/__init__.py b/torch/export/__init__.py
index f117a7b..4ffbb9b 100644
--- a/torch/export/__init__.py
+++ b/torch/export/__init__.py
@@ -334,7 +334,9 @@
_max = sys.maxsize - 1 if max is None else builtins.min(max, sys.maxsize - 1)
assert _max > _min, f"Cannot create Dim with inconsistent min={min}, max={max}"
dim = _Dim(name, (int,), {"min": _min, "max": _max})
- dim.__module__ = inspect.getmodule(inspect.stack()[1][0]).__name__ # type: ignore[union-attr]
+ dim.__module__ = getattr(
+ inspect.getmodule(inspect.stack()[1][0]), "__name__", "__main__"
+ )
return dim
