如何正确重写上下文管理器？解决autocast调用报错问题

错误原因 #

你遇到的'NoneType' object is not callable错误，是因为在no_autocast.__enter__方法里，把全局的autocast赋值成了do_nothing(*args, **kwargs)的返回值——而do_nothing是pass，返回None。后续调用autocast("cuda", dtype=torch.float16)时，本质是调用None，自然触发报错。

修复思路 #

要让替换后的autocast依然是可调用对象，并且调用后返回一个空的上下文管理器（进入、退出都不执行任何操作），这样with autocast(...)的语法才能正常运行。

修复后的代码 #

import torch
from torch import autocast
from contextlib import nullcontext

print(autocast)

class no_autocast():
    def __init__(self):
        self.original_autocast = None

    def __enter__(self):
        global autocast
        # 保存原autocast对象
        self.original_autocast = autocast
        # 替换成返回空上下文管理器的可调用对象
        autocast = lambda *args, **kwargs: nullcontext()

    def __exit__(self, exc_type, exc_value, exc_traceback):
        global autocast
        # 恢复原autocast
        autocast = self.original_autocast

with no_autocast():
    print("outer 0:", torch.is_autocast_enabled())
    print("outer 1:", torch.is_autocast_enabled())
    with autocast("cuda", dtype=torch.float16):
        def func(a, b): return a @ b
        print("inner ", torch.is_autocast_enabled())

关键修改点 #

1. 用contextlib.nullcontext()创建空上下文管理器：它的__enter__和__exit__都不执行任何操作，完美符合"禁用原有上下文管理器"的需求。
2. 替换autocast为一个lambda函数：这个lambda接受任意参数，返回空上下文管理器，保证autocast(...)的调用语法完全兼容原接口。
3. 变量名更清晰：把self.autocast改成self.original_autocast，避免和全局autocast混淆。

如果你不想依赖contextlib，也可以自己实现一个空上下文管理器类：

class NullContext:
    def __enter__(self):
        pass
    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

# 然后在__enter__里替换成：
autocast = lambda *args, **kwargs: NullContext()

内容的提问来源于stack exchange，提问作者user3322839