predict_mode

predict_mode()

Returns a scope context to be used in ‘with’ statement in which forward pass behavior is set to inference mode, without changing the recording states.

Example:

with autograd.record():
    y = model(x)
    with autograd.predict_mode():
        y = sampling(y)
    backward([y])