Repository: Lightning-AI/pytorch-lightning

Issue:
Summary not working for model on GPU with multiple inputs
**Describe the bug**
When you want a summary for a model that requires multiple input parameters for forward, then this doesn't work. You can set self.example_input_array to be a tuple and there is some code for passing this to the forward method. However, if the model is on cuda, it tries to pass to move this input directly to cuda without a check whether it is a tuple or list.

the line with the error is here: 
[pytorch-lightning/blob/master/pytorch_lightning/root_module/memory.py#L53](https://github.com/williamFalcon/pytorch-lightning/blob/7324dd902b8d071f4889ab1274a4d4dc09de9a78/pytorch_lightning/root_module/memory.py#L53)

example of how it should be checked: 
[pytorch-lightning/blob/master/pytorch_lightning/root_module/memory.py#L61](https://github.com/williamFalcon/pytorch-lightning/blob/7324dd902b8d071f4889ab1274a4d4dc09de9a78/pytorch_lightning/root_module/memory.py#L61)

**To Reproduce**
Steps to reproduce the behavior:
1. create a model that requires multiple inputs in the forward method. 
2. set self.example_input_array to be a tuple
3. run the model on GPU

**Expected behavior**
a list with all layers and the input and output shapes of these layers. 

**Screenshots**
![image](https://user-images.githubusercontent.com/6348139/69352469-4bc70180-0c7d-11ea-88c6-8056ab531c80.png)

**Desktop (please complete the following information):**
 - OS: Linux Mint 19.2
 - Browser chrome
 - Version 8.0.3904.97 (Official Build) (64-bit)



Discussion:
@victormilewski1994 good point. I guess we didn't specifically design the summarizer to take in multiple inputs. Would you be interested in submitting a PR?

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
@@ -50,20 +50,31 @@ def get_variable_sizes(self):
         input_ = self.model.example_input_array
 
         if self.model.on_gpu:
-            input_ = input_.cuda(0)
+            device = next(self.model.parameters()).get_device()
+            # test if input is a list or a tuple
+            if isinstance(input_, (list, tuple)):
+                input_ = [input_i.cuda(device) if torch.is_tensor(input_i) else input_i
+                          for input_i in input_]
+            else:
+                input_ = input_.cuda(device)
 
         if self.model.trainer.use_amp:
-            input_ = input_.half()
+            # test if it is not a list or a tuple
+            if isinstance(input_, (list, tuple)):
+                input_ = [input_i.half() if torch.is_tensor(input_i) else input_i
+                          for input_i in input_]
+            else:
+                input_ = input_.half()
 
         with torch.no_grad():
 
             for _, m in mods:
-                if type(input_) is list or type(input_) is tuple:  # pragma: no cover
+                if isinstance(input_, (list, tuple)):  # pragma: no cover
                     out = m(*input_)
                 else:
                     out = m(input_)
 
-                if type(input_) is tuple or type(input_) is list:  # pragma: no cover
+                if isinstance(input_, (list, tuple)):  # pragma: no cover
                     in_size = []
                     for x in input_:
                         if type(x) is list:
@@ -75,7 +86,7 @@ def get_variable_sizes(self):
 
                 in_sizes.append(in_size)
 
-                if type(out) is tuple or type(out) is list:  # pragma: no cover
+                if isinstance(out, (list, tuple)):  # pragma: no cover
                     out_size = np.asarray([x.size() for x in out])
                 else:
                     out_size = np.array(out.size())


Test patch: