AFAIK this is not super straightforward, the way fvcore works is that it’ll insert some hooks that will observe what an nn.Module is doing to count its FLOPS. torch.jit is under maintenance mode so even if it was possible to support hooks it’s unlikely to get prioritized and its not trivial to convert from a JIT model to a regular PyTorch model so your best bet is to profile the model without jit
1 Like