auto newobj = module->entity()->object()->create(oldobj->shape());
newobj->layout(coco::FeatureLayouts::BHWC::create(oldobj->shape()));
- if (auto oldbag = oldobj->bag())
+ if (oldobj->bag() != nullptr)
{
- assert(oldbag->module() == module);
+ using nncc::core::ADT::feature::num_elements;
+
+ // NOTE The size of bag should be at least "BxHxWxC" as "newobj" uses BHWC layout
+ const uint32_t batch = newobj->layout()->batch();
+ const uint32_t count = num_elements(newobj->layout()->shape());
+ const uint32_t bag_size = batch * count;
+
// Clone bag only when there is a backing bag for a given feature object
- auto newbag = module->entity()->bag()->create(oldbag->size());
+ auto newbag = module->entity()->bag()->create(bag_size);
newobj->bag(newbag);
}