Update benchmark_litgpt.py

Lightning-AI · Aug 21, 2024 · a88b136 · a88b136
1 parent 77b1928
commit a88b136
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/thunder/benchmarks/benchmark_litgpt.py b/thunder/benchmarks/benchmark_litgpt.py
@@ -351,6 +351,7 @@ def init_model(self):
         init_device = torch.device("meta") if self.distributed_mode in FSDP_MODES else self.device
         with init_device:
             model = GPT(self.config)
+        # TODO(crcrpar): Remove this guard once https://github.com/pytorch/ao/pull/713 is merged
         if (
             self.distributed_mode == "fsdp2"
             and self._torchao_fp8_handler._enabled