aligning code with official demo

watermark.py

@@ -92,7 +92,8 @@ class WatermarkingKirchenbauerLogitsProcessor(WatermarkingLogitsProcessor):
 class Watermarker(object):
     def __init__(self, modelname="facebook/opt-350m", window_size = 0, payload_bits = 0, logits_processor = None, *args, **kwargs):
         self.tokenizer = AutoTokenizer.from_pretrained(modelname, use_auth_token=hf_token)
-        self.model = AutoModelForCausalLM.from_pretrained(modelname, use_auth_token=hf_token).to(device)
+        self.model = AutoModelForCausalLM.from_pretrained(modelname, use_auth_token=hf_token, torch_dtype=torch.float16,
+        device_map='auto').to(device)
         self.model.eval()
         self.window_size = window_size