Update inference/generate_shards.py

inference/generate_shards.py

@@ -47,7 +47,7 @@ class ShardGenerator:
     
     def generate_placeholder_shards(
         self,
-        shard_size_mb: float = 2800,
+        shard_size_mb: float = 3010,
         tensor_dtype: torch.dtype = torch.bfloat16
     ):
         """
@@ -122,7 +122,7 @@ class ShardGenerator:
     def split_large_model(
         self,
         model_state_dict: Dict[str, torch.Tensor],
-        max_shard_size_gb: float = 2.8
+        max_shard_size_gb: float = 3.01
     ):
         """
         Split a large model into shards
@@ -287,7 +287,7 @@ def main():
     parser.add_argument(
         "--shard-size",
         type=float,
-        default=2800,
+        default=3010,
         help="Target shard size in MB"
     )