From bda8eb22ac772d87fde6d4fec50c0f516d3b97ac Mon Sep 17 00:00:00 2001
From: Eitan <eitan.kosman@gmail.com>
Date: Sat, 7 Nov 2020 21:22:24 +0200
Subject: [PATCH] ADD call method

---
 pytorch_wrapper/system.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/pytorch_wrapper/system.py b/pytorch_wrapper/system.py
index f992230..56b2d62 100644
--- a/pytorch_wrapper/system.py
+++ b/pytorch_wrapper/system.py
@@ -514,6 +514,9 @@ def _execute_method_on_multi_gpus(self,
 
         return results
 
+    def __call__(self, *args, **kwargs):
+        return self.model(*args, **kwargs)
+
 
 class _Trainer(object):