Repository: Lightning-AI/pytorch-lightning

Issue:
example_input_array dtype
Currently assumed that example_input_array dtype to be equal to model dtype. This is not necessarily correct - e.g. if input is a vector of INT.

https://github.com/PyTorchLightning/pytorch-lightning/blob/7dc58bd286b1e81ca4d293f05bddff5e93361020/pytorch_lightning/core/memory.py#L192


Discussion:
Hi, I don't understand. Does it throw an error or does it display nothing? Could you clarify?
I don't think we can very accurately define the "input shape" for anything other than tensors. 
For this reason we exclude things like dicts from the overview, because it is not very practical to visualize this in a table. 
Hi, currently the model is run with `input_` as an input. If the model expects a tensor of INTs then it will crash if floats come. I encountered this issue when pretraining an ALBERT-like model. This receives word embeddings as inputs, which have to be integers as they're going into a `nn.Embedding`
okay I see, so we should not change the dtype as given by example_input_array. 
I can't recall why I added this conversion, maybe  it was because of amp and the half conversions. I'll have a closer look, thanks for bringing it up.

as a workaround until it is fixed, cast your input to int before feeding to the embedding layer, or don't use the example_input_array.

Resolving patch:
diff --git a/pytorch_lightning/core/memory.py b/pytorch_lightning/core/memory.py
--- a/pytorch_lightning/core/memory.py
+++ b/pytorch_lightning/core/memory.py
@@ -208,7 +208,6 @@ def _forward_example_input(self) -> None:
 
         input_ = model.example_input_array
         input_ = model.transfer_batch_to_device(input_, model.device)
-        input_ = apply_to_collection(input_, torch.Tensor, lambda x: x.type(model.dtype))
 
         if trainer is not None and trainer.use_amp:
             if NATIVE_AMP_AVALAIBLE:


Test patch:
diff --git a/tests/core/test_memory.py b/tests/core/test_memory.py
--- a/tests/core/test_memory.py
+++ b/tests/core/test_memory.py
@@ -42,6 +42,19 @@ def forward(self, x, y):
         return out
 
 
+class MixedDtypeModel(LightningModule):
+    """ The parameters and inputs of this model have different dtypes. """
+
+    def __init__(self):
+        super().__init__()
+        self.embed = nn.Embedding(10, 20)   # expects dtype long as input
+        self.reduce = nn.Linear(20, 1)      # dtype: float
+        self.example_input_array = torch.tensor([[0, 2, 1], [3, 5, 3]])  # dtype: long
+
+    def forward(self, x):
+        return self.reduce(self.embed(x))
+
+
 @pytest.mark.parametrize(['mode'], [
     pytest.param(ModelSummary.MODE_FULL),
     pytest.param(ModelSummary.MODE_TOP),
@@ -59,15 +72,15 @@ def test_empty_model_summary_shapes(mode):
     pytest.param(ModelSummary.MODE_FULL),
     pytest.param(ModelSummary.MODE_TOP),
 ])
-@pytest.mark.parametrize(['device', 'dtype'], [
-    pytest.param(torch.device('cpu'), torch.double),
-    pytest.param(torch.device('cuda', 0), torch.float),
-    pytest.param(torch.device('cuda', 0), torch.float16),
+@pytest.mark.parametrize(['device'], [
+    pytest.param(torch.device('cpu')),
+    pytest.param(torch.device('cuda', 0)),
+    pytest.param(torch.device('cuda', 0)),
 ])
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU.")
-def test_linear_model_summary_shapes(device, dtype, mode):
+def test_linear_model_summary_shapes(device, mode):
     """ Test that the model summary correctly computes the input- and output shapes. """
-    model = UnorderedModel().type(dtype).to(device)
+    model = UnorderedModel().to(device)
     model.train()
     summary = model.summarize(mode=mode)
     assert summary.in_sizes == [
@@ -85,10 +98,23 @@ def test_linear_model_summary_shapes(device, dtype, mode):
         UNKNOWN_SIZE,
     ]
     assert model.training
-    assert model.dtype == dtype
     assert model.device == device
 
 
+def test_mixed_dtype_model_summary():
+    """ Test that the model summary works with models that have mixed input- and parameter dtypes. """
+    model = MixedDtypeModel()
+    summary = model.summarize()
+    assert summary.in_sizes == [
+        [2, 3],         # embed
+        [2, 3, 20],     # reduce
+    ]
+    assert summary.out_sizes == [
+        [2, 3, 20],     # embed
+        [2, 3, 1],      # reduce
+    ]
+
+
 @pytest.mark.parametrize(['mode'], [
     pytest.param(ModelSummary.MODE_FULL),
     pytest.param(ModelSummary.MODE_TOP),