torch.Tensor.map_

Tensor.map_(tensor, callable)

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

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

def callable(a, b) -> number