Update internally to reflect #2597

mridul-sahu · Orbax Authors · commit 3ca3bf274b61 · 2025-11-18T08:53:35.000-08:00
PiperOrigin-RevId: 833832607
diff --git a/checkpoint/orbax/checkpoint/_src/multihost/multislice.py b/checkpoint/orbax/checkpoint/_src/multihost/multislice.py
@@ -156,6 +156,7 @@ def get_device_memory() -> int:
       'NVIDIA H100 80GB HBM3': int(80e9),
       'NVIDIA H200': int(144e9),
       'NVIDIA B200': int(183e9),
+      'NVIDIA B300 SXM6 AC': int(275e9),
   }
   memory = hbm_memory.get(device.device_kind, None)
   if memory is None:

Original file line number	Diff line number	Diff line change
`@@ -156,6 +156,7 @@ def get_device_memory() -> int:`
`156`	`156`	`'NVIDIA H100 80GB HBM3': int(80e9),`
`157`	`157`	`'NVIDIA H200': int(144e9),`
`158`	`158`	`'NVIDIA B200': int(183e9),`
	`159`	`+ 'NVIDIA B300 SXM6 AC': int(275e9),`
`159`	`160`	`}`
`160`	`161`	`memory = hbm_memory.get(device.device_kind, None)`
`161`	`162`	`if memory is None:`