class Experiment(object):
    def __init__(self, epochs=1):
        self.epochs = epochs

    def get_epochs(self):
        return self.epochs

    def train(self):
        raise NotImplementedError

    def test(self):
        raise NotImplementedError

    def set_loss(self):
        raise NotImplementedError

    def checkpoint(self):
        raise NotImplementedError

    def save(self):
        raise NotImplementedError

    def load(self):
        raise NotImplementedError

    def is_done(self):
        raise NotImplementedError

class PyTorchExperiment(object):
    def save(self):
        pass

class Experiment(object):
    def __init__(self, model, optimizer, train_data, test_data, epochs=1):
        self.model = model
        self.optimizer = optimizer
        self.train_data = train_data
        self.test_data = test_data
        self.epochs = epochs
        self.loss = 0
        self.current_epoch = 0

    def get_epoch(self):
        return self.current_epoch

    def train(self):
        raise NotImplementedError

    def test(self):
        raise NotImplementedError

    def set_loss(self, loss):
        self.loss = loss

    def checkpoint(self):
        raise NotImplementedError

    def save(self):
        raise NotImplementedError

    def load(self):
        raise NotImplementedError

    def is_done(self):
        raise NotImplementedError

class PyTorchExperiment(object):
    def save(self):
        pass

Add necessary params to Experiment.
Add necessary params to Experiment.

  class Experiment(object):
-     def __init__(self, epochs=1):
+     def __init__(self, model, optimizer, train_data, test_data, epochs=1):
+         self.model = model
+         self.optimizer = optimizer
+         self.train_data = train_data
+         self.test_data = test_data
          self.epochs = epochs
+         self.loss = 0
+         self.current_epoch = 0
  
-     def get_epochs(self):
+     def get_epoch(self):
-         return self.epochs
+         return self.current_epoch
  
      def train(self):
          raise NotImplementedError
  
      def test(self):
          raise NotImplementedError
  
-     def set_loss(self):
+     def set_loss(self, loss):
-         raise NotImplementedError
+         self.loss = loss
  
      def checkpoint(self):
          raise NotImplementedError
  
      def save(self):
          raise NotImplementedError
  
      def load(self):
          raise NotImplementedError
  
      def is_done(self):
          raise NotImplementedError
  
  class PyTorchExperiment(object):
      def save(self):
          pass
  
Add necessary params to Experiment.
## Code Before:
class Experiment(object):
    def __init__(self, epochs=1):
        self.epochs = epochs

    def get_epochs(self):
        return self.epochs

    def train(self):
        raise NotImplementedError

    def test(self):
        raise NotImplementedError

    def set_loss(self):
        raise NotImplementedError

    def checkpoint(self):
        raise NotImplementedError

    def save(self):
        raise NotImplementedError

    def load(self):
        raise NotImplementedError

    def is_done(self):
        raise NotImplementedError

class PyTorchExperiment(object):
    def save(self):
        pass

## Instruction:
Add necessary params to Experiment.
## Code After:
class Experiment(object):
    def __init__(self, model, optimizer, train_data, test_data, epochs=1):
        self.model = model
        self.optimizer = optimizer
        self.train_data = train_data
        self.test_data = test_data
        self.epochs = epochs
        self.loss = 0
        self.current_epoch = 0

    def get_epoch(self):
        return self.current_epoch

    def train(self):
        raise NotImplementedError

    def test(self):
        raise NotImplementedError

    def set_loss(self, loss):
        self.loss = loss

    def checkpoint(self):
        raise NotImplementedError

    def save(self):
        raise NotImplementedError

    def load(self):
        raise NotImplementedError

    def is_done(self):
        raise NotImplementedError

class PyTorchExperiment(object):
    def save(self):
        pass
