Interesting comparison ! I’ve got the correct result only by using LBFGS optimizer :
lasso_pytorch = nn.Linear(X_train.shape[1], 1, bias=True)

mse_loss = nn.MSELoss(reduction='sum')
optimizer = optim.LBFGS(lasso_pytorch.parameters(), lr=1)
alpha = 0.2
n_epoch = 5000


lasso_pytorch.train()
for epoch in range(n_epoch):
    def closure():
        optimizer.zero_grad()
        outputs = lasso_pytorch(torch.FloatTensor(X_train))
        loss = 0.5 * mse_loss(outputs, torch.FloatTensor(y_train.reshape(-1,1)))

        for param in lasso_pytorch.parameters():
            loss += alpha *param.abs().sum()
        loss.backward()
        return loss
        
    optimizer.step(closure)
    
coef_pytorch = np.array(lasso_pytorch.weight.data).squeeze()
print(np.sort(np.argsort(-np.abs(coef_pytorch))[:10]))

It don’t know if it helps and would be happy to know why other optimizers don’t find the right coefs.