diff --git a/deepomatic/api/resources/network.py b/deepomatic/api/resources/network.py
index 7209838..db45d41 100644
--- a/deepomatic/api/resources/network.py
+++ b/deepomatic/api/resources/network.py
@@ -78,5 +78,5 @@ def create(self, *args, **kwargs):
     def _convert_result_to_numpy(result):
         new_result = {}
         for tensor in result['tensors']:
-            new_result[tensor['name']] = np.array(tensor['data']).reshape(tensor['dims'])
+            new_result[tensor['name']] = np.array(tensor['data']).reshape(tensor['shape'])
         return new_result
diff --git a/demo.py b/demo.py
index df92050..9dcd74b 100644
--- a/demo.py
+++ b/demo.py
@@ -374,7 +374,8 @@ def pretty_print_json(data):
 
 def display_inference_tensor(result):
     for tensor_name, numpy_array in result.items():
-        print_comment("tensor '{name}', dimensions: {dims}".format(name=tensor_name, dims='x'.join(map(str, numpy_array.shape))))
+        shape_str = 'x'.join(map(str, numpy_array.shape))
+        print_comment(f"tensor '{tensor_name}', shape: {shape_str}")
 
 
 if __name__ == '__main__':