torch.Tensor.map_

Tensor.map_(tensor, callable)

Применяет callable к каждому элементу в тензоре self и заданному tensor и сохраняет результаты в тензор self. Тензор self и заданный tensor должны быть совместимы по правилам трансляции.

Функция callable должна иметь следующий вид:

def callable(a, b) -> number