class_docstr: Implementation of the BYOL architecture.
  
  Attributes:
      backbone:
          Backbone model to extract features from images.
      num_ftrs:
          Dimension of the embedding (before the projection mlp).
      hidden_dim:
          Dimension of the hidden layer in the projection and prediction mlp.
      out_dim:
          Dimension of the output (after the projection/prediction mlp).
      m:
          Momentum for the momentum update of encoder.
class_docstr_tok: ["Implementation","of","the","BYOL","architecture",".","Attributes",":","backbone",":","Backbone","model","to","extract","features","from","images",".","num_ftrs",":","Dimension","of","the","embedding","(","before","the","projection","mlp",")",".","hidden_dim",":","Dimension","of","the","hidden","layer","in","the","projection","and","prediction","mlp",".","out_dim",":","Dimension","of","the","output","(","after","the","projection\/prediction","mlp",")",".","m",":","Momentum","for","the","momentum","update","of","encoder","."]
docstr: Forward pass through the encoder and the momentum encoder.
  
  Performs the momentum update, extracts features with the backbone and
  applies the projection (and prediciton) head to the output space. If
  x1 is None, only x0 will be processed otherwise, x0 is processed with
  the encoder and x1 with the momentum encoder.
  
  Args:
      x0:
          Tensor of shape bsz x channels x W x H.
      x1:
          Tensor of shape bsz x channels x W x H.
  
  Returns:
      The output proejction of x0 and (if x1 is not None) the output
      projection of x1.
  
  Examples:
      >>> # single input, single output
      >>> out = model._forward(x)
      >>>
      >>> # two inputs, two outputs
      >>> out0, out1 = model._forward(x0, x1)
docstr_tok: ["Forward","pass","through","the","encoder","and","the","momentum","encoder",".","Performs","the","momentum","update",",","extracts","features","with","the","backbone","and","applies","the","projection","(","and","prediciton",")","head","to","the","output","space",".","If","x1","is","None",",","only","x0","will","be","processed","otherwise",",","x0","is","processed","with","the","encoder","and","x1","with","the","momentum","encoder",".","Args",":","x0",":","Tensor","of","shape","bsz","x","channels","x","W","x","H.","x1",":","Tensor","of","shape","bsz","x","channels","x","W","x","H.","Returns",":","The","output","proejction","of","x0","and","(","if","x1","is","not","None",")","the","output","projection","of","x1",".","Examples",":",">",">",">","#","single","input",",","single","output",">",">",">","out","=","model._forward","(","x",")",">",">",">",">",">",">","#","two","inputs",",","two","outputs",">",">",">","out0",",","out1","=","model._forward","(","x0",",","x1",")"]
code: def _forward(self, x0: torch.Tensor, x1: torch.Tensor = None):
      """Forward pass through the encoder and the momentum encoder.
  
      Performs the momentum update, extracts features with the backbone and
      applies the projection (and prediciton) head to the output space. If
      x1 is None, only x0 will be processed otherwise, x0 is processed with
      the encoder and x1 with the momentum encoder.
  
      Args:
          x0:
              Tensor of shape bsz x channels x W x H.
          x1:
              Tensor of shape bsz x channels x W x H.
  
      Returns:
          The output proejction of x0 and (if x1 is not None) the output
          projection of x1.
  
      Examples:
          >>> # single input, single output
          >>> out = model._forward(x)
          >>>
          >>> # two inputs, two outputs
          >>> out0, out1 = model._forward(x0, x1)
  
      """
  
      self._momentum_update(self.m)
  
      # forward pass of first input x0
      f0 = self.backbone(x0).flatten(start_dim=1)
      z0 = self.projection_head(f0)
      out0 = self.prediction_head(z0)
  
      if x1 is None:
          return out0
  
      # forward pass of second input x1
      with torch.no_grad():
          f1 = self.momentum_backbone(x1).flatten(start_dim=1)
          out1 = self.momentum_projection_head(f1)
  
      return out0, out1
  
file_path: lightly/lightly/models/byol.py
imports: import warnings
  import torch
  import torch.nn
  from lightly.models._momentum import _MomentumEncoderMixin
  from lightly.models.modules import BYOLProjectionHead
label_desc: Use image node_id 2 for calling the BYOL obj's underlying member method code with example usage: obj._forward(x0, x1) and returns: out0,  out1, out0