coreai_opt.palettization.spec.TrainingStrategy

class coreai_opt.palettization.spec.TrainingStrategy[source]

Bases: ABC

Contract for a fake-palettize module’s training-time forward pass.

__init__()

Methods

train_forward(module, weight)

Return the training-time output for weight.
abstract train_forward(module, weight)[source]

Return the training-time output for weight.

Defines how the palettized weight behaves during training (e.g. frozen, straight-through, soft assignment).

Parameters:

  • module (_FakePalettizeImplBase)

  • weight (torch.Tensor)

Return type:

torch.Tensor