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: Get predicted labels from trained model.
docstr_tok: ["Get","predicted","labels","from","trained","model","."]
code: def predict(self, idx=None, loader=None):
      """Get predicted labels from trained model."""
      # get the index of the max probability
      probs = self.predict_proba(idx, loader)
      return probs.argmax(axis=1)
  
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 6 for calling the CNN obj's underlying member method code with example usage: obj.predict(idx, loader) and returns: probs