class_docstr: AlphaFold SWA (Stochastic Weight Averaging) module wrapper.
class_docstr_tok: ["AlphaFold","SWA","(","Stochastic","Weight","Averaging",")","module","wrapper","."]
code: def forward(self, batch):
      if not self.enabled:
          raise RuntimeError("AlphaFoldSWA is not enabled")
      return self.averaged_model(batch)
  
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 3 for calling the AlphaFoldSWA obj's underlying member method code with example usage: obj.forward(batch) and returns: self