Skip to content

[BUG] Lightning model_summary does not work with TensorDictParams #1487

@egaznep

Description

@egaznep

To Reproduce

import tensordict
import torch
import torch.nn as nn
import lightning.pytorch as pl
import lightning.pytorch.utilities.model_summary.model_summary as model_summary

class LM(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.weights = tensordict.TensorDictParams(
            tensordict.TensorDict({'weight': torch.randn(10, 10)})
        )
        self.example_input_array = torch.randn(2, 10)


    def forward(self, x):
        weight = self.weights.get('weight')
        return torch.matmul(x, weight)
    
model = LM()
model_summary.summarize(model)
Traceback (most recent call last):
  File "/home/[email protected]/Projects/ssmgan-semantic/mwe.py", line 21, in <module>
    model_summary.summarize(model)
  File "/home/[email protected]/Projects/ssmgan-semantic/myenv/lib/python3.12/site-packages/lightning/pytorch/utilities/model_summary/model_summary.py", line 565, in summarize
    return ModelSummary(lightning_module, max_depth=max_depth)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/[email protected]/Projects/ssmgan-semantic/myenv/lib/python3.12/site-packages/lightning/pytorch/utilities/model_summary/model_summary.py", line 228, in __init__
    self._layer_summary = self.summarize()
                          ^^^^^^^^^^^^^^^^
  File "/home/[email protected]/Projects/ssmgan-semantic/myenv/lib/python3.12/site-packages/lightning/pytorch/utilities/model_summary/model_summary.py", line 328, in summarize
    self._forward_example_input()
  File "/home/[email protected]/Projects/ssmgan-semantic/myenv/lib/python3.12/site-packages/lightning/pytorch/utilities/model_summary/model_summary.py", line 371, in _forward_example_input
    model(input_)
  File "/home/[email protected]/Projects/ssmgan-semantic/myenv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/[email protected]/Projects/ssmgan-semantic/myenv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1881, in _call_impl
    return inner()
           ^^^^^^^
  File "/home/[email protected]/Projects/ssmgan-semantic/myenv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1818, in inner
    args_result = hook(self, args)
                  ^^^^^^^^^^^^^^^^
  File "/home/[email protected]/Projects/ssmgan-semantic/myenv/lib/python3.12/site-packages/torch/utils/module_tracker.py", line 129, in _fw_pre_hook
    name = self._get_mod_name(mod)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/[email protected]/Projects/ssmgan-semantic/myenv/lib/python3.12/site-packages/torch/utils/module_tracker.py", line 97, in _get_mod_name
    self._get_mod_name(submod)
  File "/home/[email protected]/Projects/ssmgan-semantic/myenv/lib/python3.12/site-packages/torch/utils/module_tracker.py", line 91, in _get_mod_name
    if mod not in self._known_modules:
       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/[email protected]/Projects/ssmgan-semantic/myenv/lib/python3.12/weakref.py", line 459, in __contains__
    return wr in self.data
           ^^^^^^^^^^^^^^^
  File "/home/[email protected]/Projects/ssmgan-semantic/myenv/lib/python3.12/site-packages/tensordict/base.py", line 324, in __bool__
    raise RuntimeError("Converting a tensordict to boolean value is not permitted")
RuntimeError: Converting a tensordict to boolean value is not permitted

System info

Describe the characteristic of your environment:

  • Describe how the library was installed (pip, source, ...)
  • Python version
  • Versions of any other relevant libraries
import tensordict, numpy, sys, torch
print(tensordict.__version__, numpy.__version__, sys.version, sys.platform, torch.__version__)
0.10.0
 2.3.5 3.12.12 | packaged by conda-forge | (main, Oct 22 2025, 23:25:55) [GCC 14.3.0] linux 2.9.1+cu129

Additional context

Add any other context about the problem here.

Reason and Possible fixes

If you know or suspect the reason for this bug, paste the code lines and suggest modifications.

Checklist

  • I have checked that there is no similar issue in the repo (required)
  • I have read the documentation (required)
  • I have provided a minimal working example to reproduce the bug (required)

Metadata

Metadata

Assignees

Labels

bugSomething isn't working

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions