TylerYep / torchinfo

View model summaries in PyTorch!
MIT License
2.6k stars 122 forks source link

PyTorch Compile #277

Closed mohsenhariri closed 1 year ago

mohsenhariri commented 1 year ago

Describe the bug Not sure if it's a bug, but when I send a model to torch.compiler(), "summary" can't print out the model.

To Reproduce Steps to reproduce the behavior:

model_uncompiled = NN() model = torch.compile(model_uncompiled) model.to(device) summary(model=model, input_size=(input_size,))

Expected behavior print the model!

Additional context using google colab