I’m not sure, if torch.no_grad()  is in the right place.
As you can see in the source code 3 the previous grad property will be reset once the __exit__ method is called.
Since you are returning from your function, the grad will be enabled, if it was enabled before after the transformation is done.
Have a look at this small example:
def trans_func(x):
    print('Before', x.requires_grad)
    with torch.no_grad():
        print('After ', x.requires_grad)
        c = x * 2
        print('New calculation ', c.requires_grad)
        return c

x = torch.randn(1, 2, 3, requires_grad=True)
print('Grad enabled ', torch.is_grad_enabled())
c = trans_func(x)
c.requires_grad
print('Grad enabled ', torch.is_grad_enabled())

I’m not sure, if this is, what you would like to have.
Could you explain your use case a bit?