Update README.md

README.md

@@ -33,13 +33,13 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Load the tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("metagene-ai/METAGENE-1")
-model = AutoModelForCausalLM.from_pretrained("metagene-ai/METAGENE-1", torch_dtype=torch.bfloat16)
+model = AutoModelForCausalLM.from_pretrained("metagene-ai/METAGENE-1", torch_dtype=torch.bfloat16, device_map="auto")
 
 # Example input sequence
 input_sequence = "TCACCGTTCTACAATCCCAAGCTGGAGTCAAGCTCAACAGGGTCTTC"
 
 # Tokenize the input sequence and remove the [EOS] token for generation
-input_tokens = tokenizer.encode(input_sequence, return_tensors="pt", add_special_tokens=False)
+input_tokens = tokenizer.encode(input_sequence, return_tensors="pt", add_special_tokens=False).to(model.device)
 
 # Generate output from the model
 generated_tokens = model.generate(input_tokens, max_length=32)