Upload folder using huggingface_hub

config.json

@@ -33,7 +33,7 @@
   "qkv_weight_interleaved": true,
   "token_dropout": true,
   "torch_dtype": "float32",
-  "transformers_version": "4.55.0.dev0",
+  "transformers_version": "4.55.4",
   "use_cache": true,
   "vocab_list": null,
   "vocab_size": 33

esm_nv.py

@@ -1,4 +1,5 @@
 # coding=utf-8
+# noqa: license-check
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-Apache2
 # Copyright 2022 Meta and The HuggingFace Inc. team. All rights reserved.
@@ -137,7 +138,7 @@ class NVEsmEncoder(nn.Module):
         self.emb_layer_norm_after = transformer_engine.pytorch.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
         if config.position_embedding_type == "rotary":
             self.rotary_embeddings = RotaryPositionEmbedding(config.hidden_size // config.num_attention_heads)
-            self.te_rope_emb = self.rotary_embeddings(max_seq_len=config.max_position_embeddings).cuda()
+            self.te_rope_emb = self.rotary_embeddings(max_seq_len=config.max_position_embeddings)
         else:
             self.te_rope_emb = None
 
@@ -156,6 +157,12 @@ class NVEsmEncoder(nn.Module):
         """
         all_hidden_states = () if output_hidden_states else None
 
+        if self.te_rope_emb is not None:
+            te_rope_emb = self.te_rope_emb.to(hidden_states.device, non_blocking=True)
+            te_rope_emb = te_rope_emb[: hidden_states.shape[1]]
+        else:
+            te_rope_emb = None
+
         for layer_module in self.layers:
             if output_hidden_states:
                 all_hidden_states = (*all_hidden_states, hidden_states)
@@ -163,7 +170,7 @@ class NVEsmEncoder(nn.Module):
             hidden_states = layer_module(
                 hidden_states,
                 attention_mask,
-                rotary_pos_emb=self.te_rope_emb,
+                rotary_pos_emb=te_rope_emb,
             )
 
         hidden_states = self.emb_layer_norm_after(hidden_states)