❓ [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.