How can I jit a network whose output depends on a flag?

AlphaBetaGamma96: <p>Hi All,</p>
<p>I was wondering if it’s possible to jit a network whose output depends on a flag? For example, I have a neural network that has an internal flag <code>self.use_det</code>. The network is represented as a <code>nn.Module</code> whose <code>forward</code> method comprises of a few <code>nn.Linear</code> layers that eventually produce a batch of matrices in the shape <code>[B,N,N]</code> where <code>B</code> is the batch size and <code>N</code> is the number of input nodes in the input layer. However, the returned value from the network is determined by the state of the <code>self.use_det</code> flag. If this flag is set to False, the network a tensor of shape <code>[B,N,N]</code> but if <code>self.use_det = True</code> then it return a tensor of shape<code>[B,2]</code> (due to the use of using <code>torch.slogdet</code> on the tensor.</p>
<p>Now, the question. Is it possible to <code>jit</code> a network where the output depends on this flag? Because I tried naively applying <code>torch.jit.script(net)</code> but I get the follwoing error,</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "run_mcmc.py", line 53, in &lt;module&gt;
    net = torch.jit.script(net)
  File "~/anaconda3/lib/python3.8/site-packages/torch/jit/_script.py", line 942, in script
    return torch.jit._recursive.create_script_module(
  File "~/anaconda3/lib/python3.8/site-packages/torch/jit/_recursive.py", line 391, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "~/anaconda3/lib/python3.8/site-packages/torch/jit/_recursive.py", line 452, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
  File "~/anaconda3/lib/python3.8/site-packages/torch/jit/_recursive.py", line 335, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
RuntimeError: 
Previous return statement returned a value of type Tensor but this return statement returns a value of type Tuple[Tensor, Tensor]:
  File "~/main.py", line 55
    else:
      sign, logabsdet = self.slogdet(matrices)
      return sign, logabsdet
      ~~~~~~~~~~~~~~~~~~~~~~ &lt;--- HERE
</code></pre>
<p>The forward of the class concludes with,</p>
<pre><code class="lang-auto">    if(self.use_det):
      return matrices
    else:
      sign, logabsdet = self.slogdet(matrices)
      return sign, logabsdet
</code></pre>
<p>Is this possibe? Any help is apprecitated! Thank you! <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"></p>

ptrblck: <p>Yes, it should work as long as you don’t change the output type. In your case you could return an <code>empty</code> tensor in the <code>if</code> path as seen here:</p>
<pre><code class="lang-python">class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1, 10)
        self.fc2 = nn.Linear(1, 2)
        self.flag = True
        
    def forward(self, x):
        if self.flag:
            out = self.fc1(x)
            return out, torch.empty_like(out)
        else:
            out = self.fc2(x)
            out2 = out + 1
            return out, out2
    
# eager mode
model = MyModel()
x = torch.randn(1, 1)

out = model(x)
print(out[0].shape)
&gt; torch.Size([1, 10])

model.flag = False
out = model(x)
print(out[0].shape)
&gt; torch.Size([1, 2])

# script
model = torch.jit.script(model)
out = model(x)
print(out[0].shape)
&gt; torch.Size([1, 2])

model.flag = True
out = model(x)
print(out[0].shape)
&gt; torch.Size([1, 10])
</code></pre>

AlphaBetaGamma96: <p>Ah, great! Thank you <a class="mention" href="/u/ptrblck">@ptrblck</a>! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

AlphaBetaGamma96: <p>Hi ptrblck,</p>
<p>I’ve just implemented your solution and I get an error. I was wondering if you could take a look?</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "main.py", line 152, in &lt;module&gt;
    X, _ = sampler(burn_in)
  File "~/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "~/anaconda3/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "~/src/Samplers.py", line 83, in forward
    self.step()
  File "~/anaconda3/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "~/src/Samplers.py", line 67, in step
    log_a = self._log_pdf(xcand).detach_() - self._log_pdf(self.chains).detach_()  #calculate log acceptance probability 
  File "~/anaconda3/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "~/src/Samplers.py", line 47, in _log_pdf
    return torch.slogdet(self.network(x)[0])[1].mul(2).detach_()
  File "~/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
RuntimeError: nvrtc: error: failed to open libnvrtc-builtins.so.11.1.
  Make sure that libnvrtc-builtins.so.11.1 is installed correctly.
nvrtc compilation failed: 

#define NAN __int_as_float(0x7fffffff)
#define POS_INFINITY __int_as_float(0x7f800000)
#define NEG_INFINITY __int_as_float(0xff800000)


template&lt;typename T&gt;
__device__ T maximum(T a, T b) {
  return isnan(a) ? a : (a &gt; b ? a : b);
}

template&lt;typename T&gt;
__device__ T minimum(T a, T b) {
  return isnan(a) ? a : (a &lt; b ? a : b);
}

extern "C" __global__
void fused_tanh_add(double* t0, double* t1, double* aten_add) {
{
  double v = __ldg(t1 + (64 * (((512 * blockIdx.x + threadIdx.x) / 64) % 24) + ((512 * blockIdx.x + threadIdx.x) / 1536) * 1536) + (512 * blockIdx.x + threadIdx.x) % 64);
  double v_1 = __ldg(t0 + (64 * (((512 * blockIdx.x + threadIdx.x) / 64) % 24) + ((512 * blockIdx.x + threadIdx.x) / 1536) * 1536) + (512 * blockIdx.x + threadIdx.x) % 64);
  aten_add[(64 * (((512 * blockIdx.x + threadIdx.x) / 64) % 24) + ((512 * blockIdx.x + threadIdx.x) / 1536) * 1536) + (512 * blockIdx.x + threadIdx.x) % 64] = (tanh(v)) + v_1;
}
}
</code></pre>
<p>I assume this is more likely a non-pytorch issue?</p>
<p>I’m running Ubuntu 20.04, CUDA 11.2, Driver 460.80 and pytorch is 1.8.1+cu111 (So I assume it’s running CUDA 11.1, could that be an issue?)</p>
<p>Thank you! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

ptrblck: <p>This error seems to be related to <a href="https://github.com/pytorch/pytorch/issues/58101">this issue</a>, which was already fixed, so you might want to update to the latest release (<code>1.9.0</code>) or the nightly.</p>

AlphaBetaGamma96: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> saves the day once again! Thank you! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>