How to get the mask of dropout op in gradient computation

Hi!

I need to develop the gradient computation of dropout.

In src/relay/op/nn/nn.cc, it wrote:
// dropout returns the original tensor with dropout applied
// and a mask tensor (1.0 where element not dropped, 0.0 where dropped)
auto ret_type = TensorTypeNode::make(data->shape, data->dtype);
reporter->Assign(types[1], TupleTypeNode::make(Array({ret_type, ret_type})));

But, how can I get the mask in gradient computation of dropout?

Waiting for your answer, thank you very much!

Currently, I wrote as follow:

@register_gradient(“nn.dropout”)
def dropout_grad(orig, grad):

param = orig.attrs
rate= param.rate

result = dropout_gradient(grad, rate)

return [result]

It not works :frowning: