class_docstr: AlphaFold SWA (Stochastic Weight Averaging) module wrapper.
class_docstr_tok: ["AlphaFold","SWA","(","Stochastic","Weight","Averaging",")","module","wrapper","."]
code: def __init__(
      self, alphafold: nn.Module, enabled: bool, decay_rate: float
  ) -> None:
      super(AlphaFoldSWA, self).__init__()
      if enabled:
          self.averaged_model = torch.optim.swa_utils.AveragedModel(
              model=alphafold,
              avg_fn=swa_avg_fn(decay_rate=decay_rate),
          )
          self.enabled = True
      else:
          self.averaged_model = None
          self.enabled = False
  
file_path: apex/apex/contrib/test/openfold_triton/test_fused_adam_swa.py
imports: from itertools import chain
  import random
  import unittest
  import torch
  import torch.nn
label_desc: Use image node_id 1 to create a new AlphaFoldSWA object from inherited base classes: nn with example: obj = AlphaFoldSWA(alphafold, enabled, decay_rate)