def METHOD_NAME(self, gmod):
    # TODO - This needs major investigation. Currently, we are tagging all
    # the forward nodes as recomputable. However, torch.utils.checkpoint
    # provides a custom function to selectively recompute. We will have to
    # figure out how to tag seletively.
    unique_graph_id = next(uid)
    for node in gmod.graph.nodes:
        if node.op in ("call_function", "call_method", "call_module"):
            node.meta["recompute"] = unique_graph_id
    return gmod