If statement with tensor scalar in Pytorch

[J_Johnson]

There is a torch.where statement that is the if equivalent for tensors.

Usage:

a = torch.rand((128, 3, 64, 64))
b = torch.rand((128, 3, 64, 64))

c = torch.where(a<b, a, b)

https://pytorch.org/docs/stable/generated/torch.where.html

Alternatively, if you need to maintain the graph, you can either write custom backward functions, as described here:

How do I pass grad through torch.where

Hi guys, when using torch.where, it works as a threshold thus unable to pass the grad by torch.autograd.grad(…). Instead, is there a way to keep to same graph and pass through to next differentiable operation? For example, in a very simple code below import torch a = torch.tensor([[1.0, 2.0], [3.0, 4.0]], requires_grad=True) b = torch.tensor([[0.1],[0.9]]) c = torch.mm(a, b) # ** position2 **: grad couldn't pass back here c = torch.where(c>2, torch.tensor(1.0, requires_grad=True), torch.tens…

OR construct your logic to use functions that have backward well defined. For example:

b_larger = a<b
c = b_larger*a + (1-b_larger)*b