setup

dingguanglei · Nov 21, 2018 · baa3533 · baa3533
1 parent f2254da
commit baa3533
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/jdit/model.py b/jdit/model.py
@@ -114,7 +114,7 @@ def __call__(self, *args, **kwargs):
     def __getattr__(self, item):
         return getattr(self.model, item)
 
-    def define(self, proto_model: Module, gpu_ids: Union[list, tuple], init_method: Union[str, FunctionType, None] ,
+    def define(self, proto_model: Module, gpu_ids: Union[list, tuple], init_method: Union[str, FunctionType, None],
                show_structure: bool):
         """Define and wrap a pytorch module, according to CPU, GPU and multi-GPUs.
 
@@ -356,7 +356,7 @@ def _extract_module(self, data_parallel_model: DataParallel, extract_weights=Tru
             weights = self._fix_weights(weights)
         return model, weights
 
-    def _fix_weights(self, weights: OrderedDict):
+    def _fix_weights(self, weights: Union[dict, OrderedDict]):
         # fix params' key
         from collections import OrderedDict
         new_state_dict = OrderedDict()

diff --git a/setup.py b/setup.py
@@ -40,8 +40,8 @@
         'Intended Audience :: Developers',
         "Programming Language :: Python :: 3 :: Only",
         'Programming Language :: Python :: Implementation',
-        'Programming Language :: Python :: 3.5',
         'Programming Language :: Python :: 3.6',
+        'Programming Language :: Python :: 3.7',
     ],
     # 此项需要，否则卸载时报windows error
     zip_safe=False