How to use the skorch.NeuralNet function in skorch

neural_net_reg_criterion_text = """

    criterion : torch criterion (class, default=torch.nn.MSELoss)
      Mean squared error loss."""


def get_neural_net_reg_doc(doc):
    doc = neural_net_reg_doc_start + " " + doc.split("\n ", 4)[-1]
    pattern = re.compile(r'(\n\s+)(criterion .*\n)(\s.+){1,99}')
    start, end = pattern.search(doc).span()
    doc = doc[:start] + neural_net_reg_criterion_text + doc[end:]
    return doc


# pylint: disable=missing-docstring
class NeuralNetRegressor(NeuralNet, RegressorMixin):
    __doc__ = get_neural_net_reg_doc(NeuralNet.__doc__)

    def __init__(
            self,
            module,
            *args,
            criterion=torch.nn.MSELoss,
            **kwargs
    ):
        super(NeuralNetRegressor, self).__init__(
            module,
            *args,
            criterion=criterion,
            **kwargs
        )

the logit of probabilities with shape (batch_size, ).

    threshold : float (default=0.5)
      Probabilities above this threshold is classified as 1. ``threshold``
      is used by ``predict`` and ``predict_proba`` for classification."""


def get_neural_net_binary_clf_doc(doc):
    doc = neural_net_binary_clf_doc_start + " " + doc.split("\n ", 4)[-1]
    pattern = re.compile(r'(\n\s+)(criterion .*\n)(\s.+){1,99}')
    start, end = pattern.search(doc).span()
    doc = doc[:start] + neural_net_binary_clf_criterion_text + doc[end:]
    return doc


class NeuralNetBinaryClassifier(NeuralNet, ClassifierMixin):
    # pylint: disable=missing-docstring
    __doc__ = get_neural_net_binary_clf_doc(NeuralNet.__doc__)

    def __init__(
            self,
            module,
            *args,
            criterion=torch.nn.BCEWithLogitsLoss,
            train_split=CVSplit(5, stratified=True),
            threshold=0.5,
            **kwargs
    ):
        super().__init__(
            module,
            criterion=criterion,
            train_split=train_split,

criterion : torch criterion (class, default=torch.nn.MSELoss)
      Mean squared error loss."""


def get_neural_net_reg_doc(doc):
    doc = neural_net_reg_doc_start + " " + doc.split("\n ", 4)[-1]
    pattern = re.compile(r'(\n\s+)(criterion .*\n)(\s.+){1,99}')
    start, end = pattern.search(doc).span()
    doc = doc[:start] + neural_net_reg_criterion_text + doc[end:]
    return doc


# pylint: disable=missing-docstring
class NeuralNetRegressor(NeuralNet, RegressorMixin):
    __doc__ = get_neural_net_reg_doc(NeuralNet.__doc__)

    def __init__(
            self,
            module,
            *args,
            criterion=torch.nn.MSELoss,
            **kwargs
    ):
        super(NeuralNetRegressor, self).__init__(
            module,
            *args,
            criterion=criterion,
            **kwargs
        )

    # pylint: disable=signature-differs

threshold : float (default=0.5)
      Probabilities above this threshold is classified as 1. ``threshold``
      is used by ``predict`` and ``predict_proba`` for classification."""


def get_neural_net_binary_clf_doc(doc):
    doc = neural_net_binary_clf_doc_start + " " + doc.split("\n ", 4)[-1]
    pattern = re.compile(r'(\n\s+)(criterion .*\n)(\s.+){1,99}')
    start, end = pattern.search(doc).span()
    doc = doc[:start] + neural_net_binary_clf_criterion_text + doc[end:]
    return doc


class NeuralNetBinaryClassifier(NeuralNet, ClassifierMixin):
    # pylint: disable=missing-docstring
    __doc__ = get_neural_net_binary_clf_doc(NeuralNet.__doc__)

    def __init__(
            self,
            module,
            *args,
            criterion=torch.nn.BCEWithLogitsLoss,
            train_split=CVSplit(5, stratified=True),
            threshold=0.5,
            **kwargs
    ):
        super().__init__(
            module,
            criterion=criterion,
            train_split=train_split,
            *args,
            **kwargs

"""


def get_neural_net_clf_doc(doc):
    doc = neural_net_clf_doc_start + " " + doc.split("\n ", 4)[-1]
    pattern = re.compile(r'(\n\s+)(criterion .*\n)(\s.+){1,99}')
    start, end = pattern.search(doc).span()
    doc = doc[:start] + neural_net_clf_additional_text + doc[end:]
    doc = doc + neural_net_clf_additional_attribute
    return doc


# pylint: disable=missing-docstring
class NeuralNetClassifier(NeuralNet, ClassifierMixin):
    __doc__ = get_neural_net_clf_doc(NeuralNet.__doc__)

    def __init__(
            self,
            module,
            *args,
            criterion=torch.nn.NLLLoss,
            train_split=CVSplit(5, stratified=True),
            classes=None,
            **kwargs
    ):
        super(NeuralNetClassifier, self).__init__(
            module,
            *args,
            criterion=criterion,
            train_split=train_split,
            **kwargs

A list of class labels known to the classifier.

"""


def get_neural_net_clf_doc(doc):
    doc = neural_net_clf_doc_start + " " + doc.split("\n ", 4)[-1]
    pattern = re.compile(r'(\n\s+)(criterion .*\n)(\s.+){1,99}')
    start, end = pattern.search(doc).span()
    doc = doc[:start] + neural_net_clf_additional_text + doc[end:]
    doc = doc + neural_net_clf_additional_attribute
    return doc


# pylint: disable=missing-docstring
class NeuralNetClassifier(NeuralNet, ClassifierMixin):
    __doc__ = get_neural_net_clf_doc(NeuralNet.__doc__)

    def __init__(
            self,
            module,
            *args,
            criterion=torch.nn.NLLLoss,
            train_split=CVSplit(5, stratified=True),
            classes=None,
            **kwargs
    ):
        super(NeuralNetClassifier, self).__init__(
            module,
            *args,
            criterion=criterion,
            train_split=train_split,

def __init__(self, **kwargs):
        self.reg_coef_ = kwargs.get('reg_coef', 1e-4)

        del kwargs['reg_coef']
        skorch.NeuralNet.__init__(self, **kwargs)

import skorch
import numpy as np
import torch
from torch.autograd import Variable
from sklearn.metrics import f1_score


class Net(skorch.NeuralNet):

    def __init__(
            self,
            criterion=torch.nn.CrossEntropyLoss,
            clip=0.25,
            lr=20,
            ntokens=10000,
            *args,
            **kwargs
    ):
        self.clip = clip
        self.ntokens = ntokens
        super(Net, self).__init__(criterion=criterion, lr=lr, *args, **kwargs)

    def repackage_hidden(self, h):
        """Wraps hidden states in new Variables, to detach them from their history."""

self.G.recvfrom(self.root, [])

        if pretrain:
            return self.G.readout([self.root], pretrain=True)
        else:
            # XXX(minjie): could replace the following loop with propagate call.
            #for u, v in self.walk_list:
                #self.G.update_by_edge(u, v)
                # update local should be inside the update module
                #for i in self.T_MAX_RECUR:
                #    self.G.update_local(u)
            self.G.propagate(self.walk_list)
            return self.G.readout(pretrain=False)


class Net(skorch.NeuralNet):
    def __init__(self, **kwargs):
        self.reg_coef_ = kwargs.get('reg_coef', 1e-4)

        del kwargs['reg_coef']
        skorch.NeuralNet.__init__(self, **kwargs)

    def initialize_criterion(self):
        # Overriding this method to skip initializing criterion as we don't use it.
        pass

    def get_split_datasets(self, X, y=None, **fit_params):
        # Overriding this method to use our own dataloader to change the X
        # in signature to (train_dataset, valid_dataset)
        X_train, X_valid = X
        train = self.get_dataset(X_train, None)
        valid = self.get_dataset(X_valid, None)