class_docstr: Wraps a PyTorch CNN for the MNIST dataset within an sklearn template
  
  Defines ``.fit()``, ``.predict()``, and ``.predict_proba()`` functions. This
  template enables the PyTorch CNN to flexibly be used within the sklearn
  architecture -- meaning it can be passed into functions like
  cross_val_predict as if it were an sklearn model. The cleanlab library
  requires that all models adhere to this basic sklearn template and thus,
  this class allows a PyTorch CNN to be used in for learning with noisy
  labels among other things.
  
  Parameters
  ----------
  batch_size: int
  epochs: int
  log_interval: int
  lr: float
  momentum: float
  no_cuda: bool
  seed: int
  test_batch_size: int, default=None
  dataset: {'mnist', 'sklearn-digits'}
  loader: {'train', 'test'}
    Set to 'test' to force fit() and predict_proba() on test_set
  
  Note
  ----
  Be careful setting the ``loader`` param, it will override every other loader
  If you set this to 'test', but call .predict(loader = 'train')
  then .predict() will still predict on test!
  
  Attributes
  ----------
  batch_size: int
  epochs: int
  log_interval: int
  lr: float
  momentum: float
  no_cuda: bool
  seed: int
  test_batch_size: int, default=None
  dataset: {'mnist', 'sklearn-digits'}
  loader: {'train', 'test'}
    Set to 'test' to force fit() and predict_proba() on test_set
  
  Methods
  -------
  fit
    fits the model to data.
  predict
    get the fitted model's prediction on test data
  predict_proba
    get the fitted model's probability distribution over classes for test data
class_docstr_tok: ["Wraps","a","PyTorch","CNN","for","the","MNIST","dataset","within","an","sklearn","template","Defines","``",".fit","(",")","``",",","``",".predict","(",")","``",",","and","``",".predict_proba","(",")","``","functions",".","This","template","enables","the","PyTorch","CNN","to","flexibly","be","used","within","the","sklearn","architecture","--","meaning","it","can","be","passed","into","functions","like","cross_val_predict","as","if","it","were","an","sklearn","model",".","The","cleanlab","library","requires","that","all","models","adhere","to","this","basic","sklearn","template","and","thus",",","this","class","allows","a","PyTorch","CNN","to","be","used","in","for","learning","with","noisy","labels","among","other","things",".","Parameters","--","--","--","--","--","batch_size",":","int","epochs",":","int","log_interval",":","int","lr",":","float","momentum",":","float","no_cuda",":","bool","seed",":","int","test_batch_size",":","int",",","default=None","dataset",":","{","'mnist","'",",","'sklearn-digits","'","}","loader",":","{","'train","'",",","'test","'","}","Set","to","'test","'","to","force","fit","(",")","and","predict_proba","(",")","on","test_set","Note","--","--","Be","careful","setting","the","``","loader","``","param",",","it","will","override","every","other","loader","If","you","set","this","to","'test","'",",","but","call",".predict","(","loader","=","'train","'",")","then",".predict","(",")","will","still","predict","on","test","!","Attributes","--","--","--","--","--","batch_size",":","int","epochs",":","int","log_interval",":","int","lr",":","float","momentum",":","float","no_cuda",":","bool","seed",":","int","test_batch_size",":","int",",","default=None","dataset",":","{","'mnist","'",",","'sklearn-digits","'","}","loader",":","{","'train","'",",","'test","'","}","Set","to","'test","'","to","force","fit","(",")","and","predict_proba","(",")","on","test_set","Methods","--","--","--","-","fit","fits","the","model","to","data",".","predict","get","the","fitted","model","'s","prediction","on","test","data","predict_proba","get","the","fitted","model","'s","probability","distribution","over","classes","for","test","data"]
docstr: This function adheres to sklearn's "fit(X, y)" format for
  compatibility with scikit-learn. ** All inputs should be numpy
  arrays, not pyTorch Tensors train_idx is not X, but instead a list of
  indices for X (and y if train_labels is None). This function is a
  member of the cnn class which will handle creation of X, y from the
  train_idx via the train_loader.
docstr_tok: ["This","function","adheres","to","sklearn","'s","``","fit","(","X",",","y",")","''","format","for","compatibility","with","scikit-learn",".","*","*","All","inputs","should","be","numpy","arrays",",","not","pyTorch","Tensors","train_idx","is","not","X",",","but","instead","a","list","of","indices","for","X","(","and","y","if","train_labels","is","None",")",".","This","function","is","a","member","of","the","cnn","class","which","will","handle","creation","of","X",",","y","from","the","train_idx","via","the","train_loader","."]
code: def fit(
      self,
      train_idx,
      train_labels=None,
      sample_weight=None,
      loader="train",
  ):
      """This function adheres to sklearn's "fit(X, y)" format for
      compatibility with scikit-learn. ** All inputs should be numpy
      arrays, not pyTorch Tensors train_idx is not X, but instead a list of
      indices for X (and y if train_labels is None). This function is a
      member of the cnn class which will handle creation of X, y from the
      train_idx via the train_loader."""
      if self.loader is not None:
          loader = self.loader
      if train_labels is not None and len(train_idx) != len(
          train_labels
      ):
          raise ValueError(
              "Check that train_idx and train_labels are the same length."
          )
  
      if sample_weight is not None:  # pragma: no cover
          if len(sample_weight) != len(train_labels):
              raise ValueError(
                  "Check that train_labels and sample_weight "
                  "are the same length."
              )
          class_weight = sample_weight[
              np.unique(train_labels, return_index=True)[1]
          ]
          class_weight = torch.from_numpy(class_weight).float()
          if self.cuda:
              class_weight = class_weight.cuda()
      else:
          class_weight = None
  
      train_dataset = self.get_dataset(loader)
  
      # Use provided labels if not None o.w. use MNIST dataset training labels
      if train_labels is not None:
          # Create sparse tensor of train_labels with (-1)s for labels not
          # in train_idx. We avoid train_data[idx] because train_data may
          # very large, i.e. ImageNet
          sparse_labels = (
              np.zeros(
                  self.train_size
                  if loader == "train"
                  else self.test_size,
                  dtype=int,
              )
              - 1
          )
          sparse_labels[train_idx] = train_labels
          train_dataset.targets = sparse_labels
  
      train_loader = torch.utils.data.DataLoader(
          dataset=train_dataset,
          # sampler=SubsetRandomSampler(train_idx if train_idx is not None
          # else range(self.train_size)),
          sampler=SubsetRandomSampler(train_idx),
          batch_size=self.batch_size,
          **self.loader_kwargs,
      )
  
      optimizer = optim.SGD(
          self.model.parameters(), lr=self.lr, momentum=self.momentum
      )
  
      # Train for self.epochs epochs
      for epoch in range(1, self.epochs + 1):
          # Enable dropout and batch norm layers
          self.model.train()
          for batch_idx, (data, target) in enumerate(train_loader):
              if self.cuda:  # pragma: no cover
                  data, target = data.cuda(), target.cuda()
              data, target = Variable(data), Variable(target).long()
              optimizer.zero_grad()
              output = self.model(data)
              loss = F.nll_loss(output, target, class_weight)
              loss.backward()
              optimizer.step()
              if (
                  self.log_interval is not None
                  and batch_idx % self.log_interval == 0
              ):
                  print(
                      "TrainEpoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
                          epoch,
                          batch_idx * len(data),
                          len(train_idx),
                          100.0 * batch_idx / len(train_loader),
                          loss.item(),
                      ),
                  )
  
imports: from sklearn.base import BaseEstimator
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from torchvision import datasets, transforms
  from torch.autograd import Variable
  from torch.utils.data.sampler import SubsetRandomSampler
  import numpy
label_desc: Use image node_id 5 for calling the CNN obj's underlying member method code with example usage: obj.fit(train_idx, train_labels, sample_weight, loader) without return types