Here is my final code. However, I strongly recommend you to NOT use it; it is very slow. Please use the batch eigenvalue operations that are available starting from torch.1.2.
import torch
from torch.autograd import Function, gradcheck
from torch import from_numpy
import numpy as np

class TopEigen(Function):
    """Copmpute the top eigenvector of a matrix."""

    @staticmethod
    def flatten(bag, i):
        return tuple([item for sublist in bag for item in sublist[i]])

    @staticmethod
    def _bdiag(matrix):
        b, d = matrix.shape
        dlams = np.zeros((b, d, d))
        bag_ind = [[list(range(b)), [i]*b, [i]*b] for i in range(d)]
        index = tuple([TopEigen.flatten(bag_ind, i) for i in range(d)])
        dlams[index] = matrix.T.flatten()
        return dlams

    @staticmethod
    def forward(ctx, tensor):
        cov = tensor.detach().numpy()
        lams, vecs = np.linalg.eigh(cov)
        ctx.intrm = lams, vecs
        return from_numpy(vecs[:, :, -1, None])

    @staticmethod
    def backward(ctx, grad_output):
        lams, vecs = ctx.intrm
        output = grad_output.data.numpy()
        lams = lams[:, -1, None] - lams
        index = np.where(abs(lams) > 1e-6)
        lams[index] = 1/lams[index]
        dlams = TopEigen._bdiag(lams)
        pinv = np.matmul(vecs, np.matmul(dlams, vecs.transpose((0, 2, 1))))
        return from_numpy(np.matmul(np.matmul(output.transpose((0, 2, 1)),
                          pinv).transpose((0, 2, 1)), vecs[:, None, :, -1]))