record

record(train_mode=True)

Returns an autograd recording scope context to be used in ‘with’ statement and captures code that needs gradients to be calculated.

Note

When forwarding with train_mode=False, the corresponding backward should also use train_mode=False, otherwise gradient is undefined.

Example:

with autograd.record():
    y = model(x)
    backward([y])
metric.update(...)
optim.step(...)

Parameters:

train_mode: bool, default True

Whether the forward pass is in training or predicting mode. This controls the behavior of some layers such as Dropout, BatchNorm.