(loc)
| 10 | |
| 11 | @contextlib.contextmanager |
| 12 | def load_tensor_reader(loc): |
| 13 | global LOAD_TENSOR_READER |
| 14 | assert LOAD_TENSOR_READER is None |
| 15 | # load_tensor is an "op", and we will play merry hell on |
| 16 | # Inductor's memory planning if we return a tensor that |
| 17 | # aliases another tensor that we previously returned from |
| 18 | # an operator. So unlike standard ContentStoreReader use, |
| 19 | # we disable the cache so that you always get fresh storages |
| 20 | # (no aliasing for you!) |
| 21 | LOAD_TENSOR_READER = ContentStoreReader(loc, cache=False) |
| 22 | try: |
| 23 | yield |
| 24 | finally: |
| 25 | LOAD_TENSOR_READER = None |
| 26 | |
| 27 | |
| 28 | def register_debug_prims(): |
searching dependent graphs…