Update README.md

README.md

@@ -20,19 +20,25 @@ import torch
 from transformers import AutoTokenizer, AutoModel
 
 model_id = "heyongxin233/DETree"
-tgt_layer = 18                  # 0 = embeddings; 1..24 = encoder layers (RoBERTa-large)
+tgt_layer = 18  # 0=embeddings; 1..24=encoder layers (RoBERTa-large)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print("Using device:", device)
 
 tok = AutoTokenizer.from_pretrained(model_id)
-enc = AutoModel.from_pretrained(model_id, output_hidden_states=True)
+enc = AutoModel.from_pretrained(model_id, output_hidden_states=True).to(device)
+enc.eval()
 
 texts = ["An example sentence.", "Another one."]
 batch = tok(texts, padding=True, truncation=True, return_tensors="pt")
+batch = {k: v.to(device) for k, v in batch.items()}
 
-with torch.no_grad():
+with torch.inference_mode():
     out = enc(**batch)
-    hs = out.hidden_states[tgt_layer]              # (bsz, seq, hidden)
-    mask = batch["attention_mask"].unsqueeze(-1)   # (bsz, seq, 1)
+    hs = out.hidden_states[tgt_layer]            # (bsz, seq, hidden)
+    mask = batch["attention_mask"].unsqueeze(-1)  # (bsz, seq, 1)
     hs = hs.masked_fill(~mask.bool(), float("-inf"))
-    emb, _ = hs.max(dim=1)                         # (bsz, hidden) max-pool over tokens
+    emb, _ = hs.max(dim=1)                       # max-pool over tokens
     emb = torch.nn.functional.normalize(emb, p=2, dim=-1)
-print(emb.shape)  # -> (batch_size, 1024)
+
+print(emb.device, emb.shape)  # -> (batch_size, 1024)
+```