If statement with tensor scalar in Pytorch

FurinaWai77: <p>Suppose I have two tensor type scalars a,b. I want to compare them in a “if” statement, but when I use “if a&lt;b” then there would be a detach procedure taking around 70ms on my machine. Is there an alternative to realize this? a&lt;b is a tensor bool type and I don’t know how to combine with “if” without detach. Thank you!</p>

ptrblck: <p>The <code>if</code> statement should not detach the tensors but synchronize the code, since you have introduced data-dependent control flow. The host needs to synchronize to be able to read the actual values in <code>a</code> and <code>b</code> before evaluating the <code>if</code> condition.</p>

J_Johnson: <p>There is a <code>torch.where</code> statement that is the if equivalent for tensors.</p>
<p>Usage:</p>
<pre><code class="lang-auto">
a = torch.rand((128, 3, 64, 64))
b = torch.rand((128, 3, 64, 64))

c = torch.where(a&lt;b, a, b)
</code></pre>
<p><a href="https://pytorch.org/docs/stable/generated/torch.where.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/generated/torch.where.html</a></p>
<p>Alternatively, if you need to maintain the graph, you can either write custom backward functions, as described here:</p><aside class="quote quote-modified" data-post="1" data-topic="74671">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/themoonboy/48/13133_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-do-i-pass-grad-through-torch-where/74671">How do I pass grad through torch.where</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    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&gt;2, torch.tensor(1.0, requires_grad=True), torch.tens…
  </blockquote>
</aside>

<p>OR construct your logic to use functions that have backward well defined. For example:</p>
<pre><code class="lang-auto">b_larger = a&lt;b
c = b_larger*a + (1-b_larger)*b
</code></pre>

FurinaWai77: <p>Thank you! But what I’m going to realize after (a&lt;b) is a bit troublesome. Suppose I have a 3x2x2 tensor d1 and 5x2x2 tensor d2 and I want to realize if(a&lt;b) d1=d2. If I insert a&lt;b in the multiplication then there’ll be a dimension mismatch. I want to find a way to realize a “if” on gpus or tensor bool type varible. Thank you anyway!</p>

J_Johnson: <aside class="quote no-group" data-username="FurinaWai77" data-post="4" data-topic="194899">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/f/58f4c7/48.png" class="avatar"> FurinaWai77:</div>
<blockquote>
<p>Suppose I have a 3x2x2 tensor d1 and 5x2x2 tensor d2 and I want to realize if(a&lt;b) d1=d2. If I insert a&lt;b in the multiplication then there’ll be a dimension mismatch.</p>
</blockquote>
</aside>
<p>If the dims don’t match, then you’ll need to first establish how you want define what <code>a&lt;b</code> means. There are always methods to get the dims to match, but we need to establish what you want that inequality to mean in this case.</p>