is this module adapted for tensorflow 2.0 and above?

[sohamhans19994]

No description provided.

[duhaime]

@yaroslavvb it looks like Tensorflow 2.x removed tensorflow.contrib.graph_editor, so one needs to update the logic to access get_backward_walk_ops to something like:

from tensorflow.python.ops import op_selector

# then...
op_selector.get_backward_walk_ops(...)

However, there's no op_selector.get_forward_walk_ops...

There's good discussion of a tf2 implementation of gradient checkpointing here.