Batch-wise Matrix-Free LSQR

I am trying to solve a least square problem for a linear system Ax - b, where x is a 1D vector , b is also a 1D vector that contains two sub-vectors b1 and b2, and A contains two operators: the identity map I and the differentiation operator D. This can be written as

A = torch.cat([I,  D]), b = torch.cat([b1, b2])

where

x = b1, D(x) = b2

The operator D is written as a function that takes the 1st order derivative of the signal x. The scipy function

scipy.sparse.linalg.lsqr

with first argument a function can solve the above problem matrix-free by the linear operator

scipy.sparse.linalg.LinearOperator

I am wondering whether there are similar features for

torch.linalg.lstsq
```.
Or is there any alternative way to achieve this using pytorch?