❓ [QUESTION] Custom layer with control structure not supported?

[feifzhou]

I implemented a customized layer which does different things depending on the data

...
forward(self, data):
if data.some_criterion:
do this
else:
do that
..

However, the autograd computation of forces fails on the first step:

return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.

Is it because of JIT? Is there a way to make control structure work in a customized operator?

thank you

[ThePauliPrinciple]

A practical solution is to calculate both cases for all elements, multiply them with a mask corresponding to the condition and add the two results together.