import tensorflow.contrib.graph_editor as ge

            
        elif checkpoints == 'speed':
            # checkpoint all expensive ops to maximize running speed
            checkpoints = ge.filter_ts_from_regex(fwd_ops, 'conv2d|Conv|MatMul')
            
        elif checkpoints == 'memory':
