def _patched_computation_module(module, compute_list, compute_fn):
    """
    Patch the module to compute a module's parameters, like FLOPs.

    Calls compute_fn and appends the results to compute_list.
    """
    ty = type(module)
    typestring = module.__repr__()

    class ComputeModule(ty):
        orig_type = ty

        def _original_forward(self, *args, **kwargs):
            return ty.forward(self, *args, **kwargs)

        def forward(self, *args, **kwargs):
            out = self._original_forward(*args, **kwargs)
            compute_list.append(compute_fn(self, args[0], out))
            return out

        def __repr__(self):
            return typestring

    return ComputeModule