RuntimeError: Detected that you are using FX to torch.jit.trace a dynamo-optimized function. This is not supported at the moment.

Hi,

I'm calculating flops for open-sora model with backbone of STDiT and three inputs: x, timestamp, and y, like the following:

from fvcore.nn import FlopCountAnalysis
flops = FlopCountAnalysis(model, (x, torch.randint(0, 10, (x.shape[0],), device=x.device), model_args['y']))
flops.total()

I get this error: "RuntimeError: Detected that you are using FX to torch.jit.trace a dynamo-optimized function. This is not supported at the moment."
Could you please advise if there is any solution to resolve this? Thanks