PCGrad

class torchjd.aggregation.PCGrad

GramianWeightedAggregator as defined in Algorithm 1 of Gradient Surgery for Multi-Task Learning.

__call__(matrix, /)

Computes the aggregation from the input matrix and applies all registered hooks.

Parameters:

matrix (Tensor) – The Jacobian to aggregate.

Return type:

Tensor

class torchjd.aggregation.PCGradWeighting

Weighting [PSDMatrix] giving the weights of PCGrad.

__call__(gramian, /)

Computes the vector of weights from the input Gramian and applies all registered hooks.

Parameters:

gramian (Tensor) – The Gramian from which the weights must be extracted.

Return type:

Tensor