I cannot really see how it can be a worse experience than waiting for the model to run, fix the error found, wait a second time, fix another error; when you can just check instantly if the model architecture is sound.
* Fire up Jupyter notebook, create some PyTorch nn.Module, validate the output is expected.
* Built out the model, validate the output is expected.
* Check the data shape, do a eval() run with the data.
* Move out of Jupyter notebook (or not), to actually run the trainer.
Each of the "build out model" phase, shape inference at compile time would be kind of annoying especially if you have a LSP constantly running compilation.