From 75fc4e5ff8d99c7b4622a2e99ab46f4d2c76e77b Mon Sep 17 00:00:00 2001 From: 51616 Date: Fri, 30 May 2025 08:11:31 +0000 Subject: [PATCH] autocast bfloat16 --- src/ctx_to_lora/modeling/hypernet.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/ctx_to_lora/modeling/hypernet.py b/src/ctx_to_lora/modeling/hypernet.py index 1f0f14e..7adf258 100644 --- a/src/ctx_to_lora/modeling/hypernet.py +++ b/src/ctx_to_lora/modeling/hypernet.py @@ -437,7 +437,7 @@ class HyperLoRA(nn.Module): # hidden_size=self.config.base_hidden_size, # ) - # @torch.autocast(device_type="cuda", dtype=torch.bfloat16) + @torch.autocast(device_type="cuda", dtype=torch.bfloat16) def _to_lora_dict( self, flat_loras: Float[Tensor, "bs n_layers n_modules r max_io_dim"] ) -> dict[str, dict[str, Float[Tensor, "bs n_layers r _"]]]: @@ -491,7 +491,7 @@ class HyperLoRA(nn.Module): ) return {k: v for k, v in zip(self.extra_modules, layernorms)} - # @torch.autocast(device_type="cuda", dtype=torch.bfloat16) + @torch.autocast(device_type="cuda", dtype=torch.bfloat16) def forward( self, features: Float[Tensor, "bs seq_len feature_dim"], @@ -521,7 +521,7 @@ class HyperLoRA(nn.Module): return flat_loras, flat_layernorms - # @torch.autocast(device_type="cuda", dtype=torch.bfloat16) + @torch.autocast(device_type="cuda", dtype=torch.bfloat16) def generate_weights( self, features: Float[Tensor, "bs seq_len feature_dim"],