GradientNormClippingTrainingCallback

class GradientNormClippingTrainingCallback(max_norm: float, norm_type: float | None = None)

Bases: TrainingCallback

A callback for gradient clipping before stepping the optimizer with torch.nn.utils.clip_grad_norm_().

Initialize the callback.

Parameters:

• max_norm (float) – The maximum gradient norm for use with gradient clipping.

• norm_type (float | None) – The gradient norm type to use for maximum gradient norm, cf. torch.nn.utils.clip_grad_norm_()

Methods Summary

pre_step(**kwargs)

Call before the optimizer's step.

Methods Documentation

pre_step(**kwargs: Any) → None

Call before the optimizer’s step.

Parameters:

kwargs (Any)

Return type:

None