class_docstr: Collator for IJEPA model [0].
  
  Experimental: Support for I-JEPA is experimental, there might be breaking changes
  in the future.
  
  Code inspired by [1].
  
  - [0]: Joint-Embedding Predictive Architecture, 2023, https://arxiv.org/abs/2301.08243
  - [1]: https://github.com/facebookresearch/ijepa
class_docstr_tok: ["Collator","for","IJEPA","model","[","0","]",".","Experimental",":","Support","for","I-JEPA","is","experimental",",","there","might","be","breaking","changes","in","the","future",".","Code","inspired","by","[","1","]",".","-","[","0","]",":","Joint-Embedding","Predictive","Architecture",",","2023",",","https",":","\/\/arxiv.org\/abs\/2301.08243","-","[","1","]",":","https",":","\/\/github.com\/facebookresearch\/ijepa"]
docstr: Create encoder and predictor masks when collating imgs into a batch
  # 1. sample enc block (size + location) using seed
  # 2. sample pred block (size) using seed
  # 3. sample several enc block locations for each image (w/o seed)
  # 4. sample several pred block locations for each image (w/o seed)
  # 5. return enc mask and pred mask
docstr_tok: ["Create","encoder","and","predictor","masks","when","collating","imgs","into","a","batch","#","1.","sample","enc","block","(","size","+","location",")","using","seed","#","2.","sample","pred","block","(","size",")","using","seed","#","3.","sample","several","enc","block","locations","for","each","image","(","w\/o","seed",")","#","4.","sample","several","pred","block","locations","for","each","image","(","w\/o","seed",")","#","5.","return","enc","mask","and","pred","mask"]
code: def __call__(self, batch):
      """
      Create encoder and predictor masks when collating imgs into a batch
      # 1. sample enc block (size + location) using seed
      # 2. sample pred block (size) using seed
      # 3. sample several enc block locations for each image (w/o seed)
      # 4. sample several pred block locations for each image (w/o seed)
      # 5. return enc mask and pred mask
      """
      B = len(batch)
  
      collated_batch = torch.utils.data.default_collate(batch)
  
      seed = self.step()
      g = torch.Generator()
      g.manual_seed(seed)
      p_size = self._sample_block_size(
          generator=g,
          scale=self.pred_mask_scale,
          aspect_ratio_scale=self.aspect_ratio,
      )
      e_size = self._sample_block_size(
          generator=g,
          scale=self.enc_mask_scale,
          aspect_ratio_scale=(1.0, 1.0),
      )
  
      collated_masks_pred, collated_masks_enc = [], []
      min_keep_pred = self.height * self.width
      min_keep_enc = self.height * self.width
      for _ in range(B):
          masks_p, masks_C = [], []
          for _ in range(self.npred):
              mask, mask_C = self._sample_block_mask(p_size)
              masks_p.append(mask)
              masks_C.append(mask_C)
              min_keep_pred = min(min_keep_pred, len(mask))
          collated_masks_pred.append(masks_p)
  
          acceptable_regions = masks_C
  
          if self.allow_overlap:
              acceptable_regions = None
  
          masks_e = []
          for _ in range(self.nenc):
              mask, _ = self._sample_block_mask(
                  e_size, acceptable_regions=acceptable_regions
              )
              masks_e.append(mask)
              min_keep_enc = min(min_keep_enc, len(mask))
          collated_masks_enc.append(masks_e)
  
      collated_masks_pred = [
          [cm[:min_keep_pred] for cm in cm_list]
          for cm_list in collated_masks_pred
      ]
      collated_masks_pred = torch.utils.data.default_collate(
          collated_masks_pred
      )
      # --
      collated_masks_enc = [
          [cm[:min_keep_enc] for cm in cm_list]
          for cm_list in collated_masks_enc
      ]
      collated_masks_enc = torch.utils.data.default_collate(
          collated_masks_enc
      )
  
      return collated_batch, collated_masks_enc, collated_masks_pred
  
imports: import math
  from multiprocessing import Value
  from typing import List, Optional, Tuple, Union
  from warnings import warn
  import torch
  import torch.nn
  import torchvision
  import torchvision.transforms
  from PIL import Image
  from lightly.transforms import GaussianBlur, Jigsaw, RandomSolarization
  from lightly.transforms.random_crop_and_flip_with_grid import RandomResizedCropAndFlip
  from lightly.transforms.rotation import random_rotation_transform
  from lightly.transforms.utils import IMAGENET_NORMALIZE
label_desc: Use image node_id 5 for calling the IJEPAMaskCollator obj's underlying member method code with example usage: obj.__call__(batch) and returns: collated_batch,  collated_masks_enc,  collated_masks_pred