fix(//core/conversion/converters/impl/linear): In inserting flatten for

layers other than conv, did not remove batch size from tensor shape Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
pytorch · Mar 28, 2020 · 377ad67 · 377ad67
1 parent 36b2eb4
commit 377ad67
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/core/conversion/converters/impl/linear.cpp b/core/conversion/converters/impl/linear.cpp
@@ -27,7 +27,7 @@ auto linear_registrations = RegisterNodeConversionPatterns()
                 new_shape.push_back(shape[0]);
                 new_shape.push_back(1);
                 new_shape.push_back(1);
-                new_shape.push_back(util::volume(util::toDims(shape)));
+                new_shape.push_back(util::volume(util::toDims(shape)) / shape[0]);
 
                 auto new_dims = util::toDims(new_shape);
                 LOG_DEBUG("Input shape is less than 4D got: " << util::toDims(shape) << ", inserting shuffle layer to reshape to 4D tensor shape: " << new_dims);