Add unit test for summary_string

Author: greenmonn

.gitignore

@@ -1,2 +1,3 @@
 __pycache__
 *.pyc
+.vscode/

torchsummary/__init__.py

@@ -1 +1 @@
-from .torchsummary import summary
+from .torchsummary import summary, summary_string

torchsummary/tests/unit_tests/torchsummary_test.py

@@ -1,8 +1,10 @@
 import unittest
-from torchsummary import summary
+from torchsummary import summary, summary_string
 from torchsummary.tests.test_models.test_model import SingleInputNet, MultipleInputNet, MultipleInputNetDifferentDtypes
 import torch
 
+gpu_if_available = "cuda:0" if torch.cuda.is_available() else "cpu"
+
 class torchsummaryTests(unittest.TestCase):
     def test_single_input(self):
         model = SingleInputNet()
@@ -15,7 +17,8 @@ def test_multiple_input(self):
         model = MultipleInputNet()
         input1 = (1, 300)
         input2 = (1, 300)
-        total_params, trainable_params = summary(model, [input1, input2], device="cpu")
+        total_params, trainable_params = summary(
+            model, [input1, input2], device="cpu")
         self.assertEqual(total_params, 31120)
         self.assertEqual(trainable_params, 31120)
 
@@ -28,9 +31,10 @@ def test_single_layer_network(self):
 
     def test_single_layer_network_on_gpu(self):
         model = torch.nn.Linear(2, 5)
-        model.cuda()
+        if torch.cuda.is_available():
+            model.cuda()
         input = (1, 2)
-        total_params, trainable_params = summary(model, input, device="cuda:0")
+        total_params, trainable_params = summary(model, input, device=gpu_if_available)
         self.assertEqual(total_params, 15)
         self.assertEqual(trainable_params, 15)
 
@@ -39,9 +43,22 @@ def test_multiple_input_types(self):
         input1 = (1, 300)
         input2 = (1, 300)
         dtypes = [torch.FloatTensor, torch.LongTensor]
-        total_params, trainable_params = summary(model, [input1, input2], device="cpu", dtypes=dtypes)
+        total_params, trainable_params = summary(
+            model, [input1, input2], device="cpu", dtypes=dtypes)
         self.assertEqual(total_params, 31120)
         self.assertEqual(trainable_params, 31120)
 
+
+class torchsummarystringTests(unittest.TestCase):
+    def test_single_input(self):
+        model = SingleInputNet()
+        input = (1, 28, 28)
+        result, (total_params, trainable_params) = summary_string(
+            model, input, device="cpu")
+        self.assertEqual(type(result), str)
+        self.assertEqual(total_params, 21840)
+        self.assertEqual(trainable_params, 21840)
+
+
 if __name__ == '__main__':
     unittest.main(buffer=True)