Как замаскировать тензор, не теряя градиент?

у меня есть тензор

import torch
a = torch.randn(1, 3, requires_grad=True)
print('a: ', a)
>>> a:  tensor([[0.0200, 1.00200, -4.2000]], requires_grad=True)

И маска

mask = torch.zeros_like(a)
mask[0][0] = 1

Я хочу замаскировать свой тензор a, не распространяя градиенты на тензор маски (в моем реальном случае у него есть градиент). Я попробовал следующее

with torch.no_grad():
    b = a * mask
    print('b: ', b)
    >>> b:  tensor([[0.0200, 0.0000, -0.0000]])

Но он полностью удаляет градиент из моего тензора. Как правильно это сделать?