| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| checkpoint = "HuggingFaceTB/SmolLM2-135M" | |
| # "cpu" for CPU usage | |
| # "gpu" if Torch is compiled with CUDA enabled | |
| device = "cpu" | |
| tokenizer = AutoTokenizer.from_pretrained(checkpoint) | |
| model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device) | |
| inputs = tokenizer.encode("Gravity is", return_tensors="pt").to(device) | |
| outputs = model.generate(inputs) | |
| print(tokenizer.decode(outputs[0])) | |