Fix fn.coord_transform handling of a default matrix in variable batch…

… case (#3958) * Fix fn.coord_transform handling of default matrix in variable batch case Signed-off-by: Kamil Tokarski <[email protected]>
NVIDIA · Jun 20, 2022 · 9c519c6 · 9c519c6
1 parent 0e27abe
commit 9c519c6
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/dali/operators/geometry/mt_transform_attr.cc b/dali/operators/geometry/mt_transform_attr.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+// Copyright (c) 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -127,11 +127,10 @@ void MTTransformAttr::ProcessMatrixArg(const OpSpec &spec, const ArgumentWorkspa
     if (static_cast<int>(mtx_.size()) != output_pt_dim_ * input_pt_dim_) {
       mtx_.resize(output_pt_dim_ * input_pt_dim_, 0);
       MakeDiagonalMatrix(mtx_, 1);
-      Repeat(per_sample_mtx_, mtx_, N);
-      if (is_fused) {
-        translation_.resize(output_pt_dim_, 0);
-        Repeat(per_sample_translation_, translation_, N);
-      }
+    }
+    Repeat(per_sample_mtx_, mtx_, N);
+    if (is_fused) {
+      per_sample_translation_.resize(output_pt_dim_ * N, 0);
     }
   }
 }

diff --git a/dali/test/python/test_dali_variable_batch_size.py b/dali/test/python/test_dali_variable_batch_size.py
@@ -291,6 +291,8 @@ def numba_setup_out_shape(out_shape, in_shape):
     (fn.cast, {'dtype': types.INT32}),
     (fn.color_space_conversion, {'image_type': types.BGR, 'output_type': types.RGB}),
     (fn.coord_transform, {'M': .5, 'T': 2}),
+    (fn.coord_transform, {'T': 2}),
+    (fn.coord_transform, {'M': .5}),
     (fn.crop, {'crop': (5, 5)}),
     (fn.erase, {'anchor': [0.3], 'axis_names': "H", 'normalized_anchor': True,
                 'shape': [0.1], 'normalized_shape': True}),