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"]
code: def __init__(
      self,
      batch_size=64,
      epochs=6,
      log_interval=50,  # Set to None to not print
      lr=0.01,
      momentum=0.5,
      no_cuda=False,
      seed=1,
      test_batch_size=None,
      dataset="mnist",
      loader=None,
  ):
      self.batch_size = batch_size
      self.epochs = epochs
      self.log_interval = log_interval
      self.lr = lr
      self.momentum = momentum
      self.no_cuda = no_cuda
      self.seed = seed
      self.cuda = not self.no_cuda and torch.cuda.is_available()
      torch.manual_seed(self.seed)
      if self.cuda:  # pragma: no cover
          torch.cuda.manual_seed(self.seed)
  
      # Instantiate PyTorch model
      self.model = SimpleNet()
      if self.cuda:  # pragma: no cover
          self.model.cuda()
  
      self.loader_kwargs = (
          {"num_workers": 1, "pin_memory": True} if self.cuda else {}
      )
      self.loader = loader
      self._set_dataset(dataset)
      if test_batch_size is not None:
          self.test_batch_size = test_batch_size
      else:
          self.test_batch_size = self.test_size
  
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 1 to create a new CNN object from inherited base classes: BaseEstimator with example: obj = CNN(batch_size, epochs, log_interval, lr, momentum, no_cuda, seed, test_batch_size, dataset, loader)