fix sharded model flag

Author: Ashish Gupta

src/sagemaker/serve/builder/model_builder.py

@@ -1293,6 +1293,9 @@ def _model_builder_optimize_wrapper(
                 max_runtime_in_sec=max_runtime_in_sec,
             )
 
+        if sharding_config:
+            self.pysdk_model._is_sharded_model = True
+
         if input_args:
             self.sagemaker_session.sagemaker_client.create_optimization_job(**input_args)
             job_status = self.sagemaker_session.wait_for_optimization_job(job_name)
@@ -1302,9 +1305,6 @@ def _model_builder_optimize_wrapper(
         if not speculative_decoding_config:
             self.pysdk_model.remove_tag_with_key(Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER)
 
-        if sharding_config:
-            self.pysdk_model._is_sharded_model = True
-
         return self.pysdk_model
 
     def _optimize_for_hf(