I don’t think the slowness is due to the python binding since I tested the following c++ code
#include <torch/extension.h>
/// param x: (batch, sample, feature)
/// param y: (batch, sample)
/// return (batch, feature)
at::Tensor lr_batch_fit(const at::Tensor &x, const at::Tensor &y){
    const int64_t batch = x.size(0), sample = x.size(1), features = x.size(2);

    const at::Tensor & xt = at::transpose(x, 1, 2);
    const at::Tensor & s = at::bmm(xt, x);
    const at::Tensor & yp = y.view({batch, sample, 1});
    const at::Tensor & pre_beta = at::bmm(xt, yp);

    at::Tensor beta = at::zeros({batch, features, 1}, x.options());
    return beta.squeeze(2);

and the corresponding python version
import torch

def lr_batch_fit_py(x, y):
    """
    :param x: (batch, sample, features)
    :param y: (batch, sample,)
    :return: beta (batch, features)
    """
    batch, sample, features = x.shape
    xt = torch.transpose(x, 1, 2)
    s = torch.bmm(xt, x)
    y = y.view(batch, sample, 1)
    pre_beta = torch.bmm(xt, y)
    beta = torch.zeros(batch, features, 1, dtype=x.dtype, device=x.device)
    return beta.squeeze(2)

both versions have mildly difference and c++ wins the race.