It raised TypeError: Type Join Failed: dtype1 = Float32, dtype2 = Int64. in Graph mode. No problem in Pynative mode.

import numpy as np
import mindspore as ms
import mindspore.ops as ops
from mindspore import Tensor, nn, context

class Graph(nn.Cell):
    def __init__(self):
        super().__init__()
        self.net = nn.Dense(10, 1)
    def construct(self, x, y):
        logits = self.net(x)
        extra_output = ops.nonzero(y > 0)
        return logits, extra_output

def main():
    context.set_context(mode=context.GRAPH_MODE)
    net = Graph()
    x = Tensor(np.random.randn(16, 10), dtype=ms.float32)
    y = Tensor([0, 0, 1, 1])
    labels = Tensor(np.random.randn(16, 1), dtype=ms.float32)
    logit, extra_output = net(x, y)
    assert extra_output.shape == (2, 1)
    loss_fn = nn.MSELoss()
    def forward(x, y, labels):
        logits, extra_output = net(x, y)
        loss = loss_fn(logits, labels)
        return loss, logits, extra_output
    grad_fn = ops.value_and_grad(forward, grad_position=None, weights=net.trainable_params(), has_aux=True)
    # raised TypeError here
    (loss, logits, extra_output), inputs_gradient = grad_fn(x, y, labels)
    assert extra_output.shape == (2, 1)

if __name__ == '__main__':
    main()