rohannsinghal
/

nemotron-custom-handler

Model card Files Files and versions

rohannsinghal commited on 28 days ago

Commit

98b9cc8

·

verified ·

1 Parent(s): 1d35377

Update handler.py

Files changed (1) hide show

handler.py +1 -1

handler.py CHANGED Viewed

@@ -4,7 +4,7 @@ class EndpointHandler:
     def __init__(self, path=""):
         # Here is the magic override that bypasses the Hugging Face bug
         print("Initializing Nemotron 8B with trust_remote_code=True...")
-        self.model = SentenceTransformer("nvidia/llama-embed-nemotron-8b", trust_remote_code=True)
         print("Model loaded successfully!")
     def __call__(self, data):

     def __init__(self, path=""):
         # Here is the magic override that bypasses the Hugging Face bug
         print("Initializing Nemotron 8B with trust_remote_code=True...")
+        self.model = SentenceTransformer("nvidia/llama-embed-nemotron-8b", trust_remote_code=True,model_kwargs={"attn_implementation": "eager"})
         print("Model loaded successfully!")
     def __call__(self, data):