Can I use pytorch-like training process instead of scikit-compatible way of use this lib?

[kaiwang0112006]

I want to do research on tabnet with federated learning, which means I need to get the model weight out and set it back during each epoch of training. It would be easier with a pytorch-like epoch training process of using this lib instead of scikit-compatible way of training.