Updated graph tracing function.

Author: gmalivenko

pytorch2keras/converter.py

@@ -82,7 +82,7 @@ def pytorch_to_keras(
     orig_state_dict_keys = _unique_state_dict(model).keys()
 
     with set_training(model, training):
-        trace, torch_out = torch.jit.trace(model, args)
+        trace, torch_out = torch.jit.get_trace_graph(model, args)
 
     if orig_state_dict_keys != _unique_state_dict(model).keys():
         raise RuntimeError("state_dict changed after running the tracer; "

pytorch2keras/layers.py

@@ -660,7 +660,8 @@ def convert_reduce_sum(params, w_name, scope_name, inputs, layers, weights):
     print('Converting reduce_sum ...')
 
     keepdims = params['keepdims'] > 0
-    target_layer = lambda x: keras.backend.sum(x, keepdims=keepdims, axis=params['axes'])
+    axis = np.array(params['axes'])
+    target_layer = lambda x: keras.backend.sum(x, keepdims=keepdims, axis=axis)
 
     lambda_layer = keras.layers.Lambda(target_layer)
     layers[scope_name] = lambda_layer(layers[inputs[0]])