import numpy as np
from .._serializable import Serializable, Serializer, Deserializer


class Model(Serializable):
    """ This is the superclass of all models.
    """

    def __init__(self, model=None):
        """ Wrap a callable model as a SHAP Model object.
        """
        if isinstance(model, Model):
            self.inner_model = model.inner_model
        else:
            self.inner_model = model

        if hasattr(model, "output_names"):
            self.output_names = model.output_names

    def __call__(self, *args):
        return np.array(self.inner_model(*args))

    def save(self, out_file):
        """ Save the model to the given file stream.
        """
        super().save(out_file)
        with Serializer(out_file, "shap.Model", version=0) as s:
            s.save("model", self.inner_model)

    @classmethod
    def load(cls, in_file, instantiate=True):
        if instantiate:
            return cls._instantiated_load(in_file)

        kwargs = super().load(in_file, instantiate=False)
        with Deserializer(in_file, "shap.Model", min_version=0, max_version=0) as s:
            kwargs["model"] = s.load("model")
        return kwargs

import numpy as np
from .._serializable import Serializable, Serializer, Deserializer
from torch import Tensor


class Model(Serializable):
    """ This is the superclass of all models.
    """

    def __init__(self, model=None):
        """ Wrap a callable model as a SHAP Model object.
        """
        if isinstance(model, Model):
            self.inner_model = model.inner_model
        else:
            self.inner_model = model

        if hasattr(model, "output_names"):
            self.output_names = model.output_names

    def __call__(self, *args):
        out = self.inner_model(*args)
        out = out.cpu().detach().numpy() if isinstance(out, Tensor) else np.array(out)
        return out

    def save(self, out_file):
        """ Save the model to the given file stream.
        """
        super().save(out_file)
        with Serializer(out_file, "shap.Model", version=0) as s:
            s.save("model", self.inner_model)

    @classmethod
    def load(cls, in_file, instantiate=True):
        if instantiate:
            return cls._instantiated_load(in_file)

        kwargs = super().load(in_file, instantiate=False)
        with Deserializer(in_file, "shap.Model", min_version=0, max_version=0) as s:
            kwargs["model"] = s.load("model")
        return kwargs

Check SHAP Model call type
Check SHAP Model call type

  import numpy as np
  from .._serializable import Serializable, Serializer, Deserializer
+ from torch import Tensor
  
  
  class Model(Serializable):
      """ This is the superclass of all models.
      """
  
      def __init__(self, model=None):
          """ Wrap a callable model as a SHAP Model object.
          """
          if isinstance(model, Model):
              self.inner_model = model.inner_model
          else:
              self.inner_model = model
  
          if hasattr(model, "output_names"):
              self.output_names = model.output_names
  
      def __call__(self, *args):
-         return np.array(self.inner_model(*args))
+         out = self.inner_model(*args)
+         out = out.cpu().detach().numpy() if isinstance(out, Tensor) else np.array(out)
+         return out
  
      def save(self, out_file):
          """ Save the model to the given file stream.
          """
          super().save(out_file)
          with Serializer(out_file, "shap.Model", version=0) as s:
              s.save("model", self.inner_model)
  
      @classmethod
      def load(cls, in_file, instantiate=True):
          if instantiate:
              return cls._instantiated_load(in_file)
  
          kwargs = super().load(in_file, instantiate=False)
          with Deserializer(in_file, "shap.Model", min_version=0, max_version=0) as s:
              kwargs["model"] = s.load("model")
          return kwargs
  
Check SHAP Model call type
## Code Before:
import numpy as np
from .._serializable import Serializable, Serializer, Deserializer


class Model(Serializable):
    """ This is the superclass of all models.
    """

    def __init__(self, model=None):
        """ Wrap a callable model as a SHAP Model object.
        """
        if isinstance(model, Model):
            self.inner_model = model.inner_model
        else:
            self.inner_model = model

        if hasattr(model, "output_names"):
            self.output_names = model.output_names

    def __call__(self, *args):
        return np.array(self.inner_model(*args))

    def save(self, out_file):
        """ Save the model to the given file stream.
        """
        super().save(out_file)
        with Serializer(out_file, "shap.Model", version=0) as s:
            s.save("model", self.inner_model)

    @classmethod
    def load(cls, in_file, instantiate=True):
        if instantiate:
            return cls._instantiated_load(in_file)

        kwargs = super().load(in_file, instantiate=False)
        with Deserializer(in_file, "shap.Model", min_version=0, max_version=0) as s:
            kwargs["model"] = s.load("model")
        return kwargs

## Instruction:
Check SHAP Model call type
## Code After:
import numpy as np
from .._serializable import Serializable, Serializer, Deserializer
from torch import Tensor


class Model(Serializable):
    """ This is the superclass of all models.
    """

    def __init__(self, model=None):
        """ Wrap a callable model as a SHAP Model object.
        """
        if isinstance(model, Model):
            self.inner_model = model.inner_model
        else:
            self.inner_model = model

        if hasattr(model, "output_names"):
            self.output_names = model.output_names

    def __call__(self, *args):
        out = self.inner_model(*args)
        out = out.cpu().detach().numpy() if isinstance(out, Tensor) else np.array(out)
        return out

    def save(self, out_file):
        """ Save the model to the given file stream.
        """
        super().save(out_file)
        with Serializer(out_file, "shap.Model", version=0) as s:
            s.save("model", self.inner_model)

    @classmethod
    def load(cls, in_file, instantiate=True):
        if instantiate:
            return cls._instantiated_load(in_file)

        kwargs = super().load(in_file, instantiate=False)
        with Deserializer(in_file, "shap.Model", min_version=0, max_version=0) as s:
            kwargs["model"] = s.load("model")
        return kwargs
