tfm.core.train_utils.try_count_flops

Counts and returns model FLOPs.

model A model instance.
inputs_kwargs An optional dictionary of argument pairs specifying inputs' shape specifications to getting corresponding concrete function.
output_path A file path to write the profiling results to.

The model's FLOPs.