Return the output of the wrapped PyTorch model for the given input,
        along with a callback to handle the backward pass.

def begin_update(self, x_data, h_0=None, drop=0.0):
        """Return the output of the wrapped PyTorch model for the given input,
        along with a callback to handle the backward pass.
        """
        x_var = torch.autograd.Variable(xp2torch(x_data), requires_grad=True)

        # Make prediction
        out, h_n = self._model(x_var, h_0)
        # Shapes will be:
        # out = seq_len, batch, hidden_size * num_directions
        # h_n = num_layers * num_directions, batch, hidden_size

        def backward_pytorch_rnn(d_data, sgd=None):
            dy_data, _ = d_data
            dout = xp2torch(dy_data)
            torch.autograd.backward((out,), grad_tensors=(dout,))
            if sgd is not None:
                if self._optimizer is None:
                    self._optimizer = self._create_optimizer(sgd)
                self._optimizer.step()
                self._optimizer.zero_grad()
            return torch2xp(x_var.grad)

        return (torch2xp(out), h_n), backward_pytorch_rnn