test_nn.py 文件源码

def default_tensor_type(type):
    type_str = torch.typename(type)
    def decorator(fn):
        @wraps(fn)
        def wrapper(*args, **kwargs):
            old_type = torch.typename(torch.Tensor())
            torch.set_default_tensor_type(type_str)
            try:
                return fn(*args, **kwargs)
            finally:
                torch.set_default_tensor_type(old_type)
        return wrapper
    return decorator