Method forward

numpy_ml/tests/nn_torch_models.py:236–246

(self, X)

        self.layer1.bias = nn.Parameter(torch.FloatTensor(intercept))

    def forward(self, X):
        # (N, H, W, C) -> (N, C, H, W)
        if X.ndim == 4:
            X = np.moveaxis(X, [0, 1, 2, 3], [0, -2, -1, -3])

        if not isinstance(X, torch.Tensor):
            X = torchify(X)

        self.X = X
        self.Y = self.layer1(self.X)
        self.Y.retain_grad()

    def extract_grads(self, X, Y_true=None):
        self.forward(X)