doc-to-lora/tmp/test_pad.py
2025-09-28 15:46:47 +00:00

8 lines
251 B
Python

import torch
from transformers import AutoTokenizer
tk = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")
x = [torch.rand(5, 3), torch.rand(10, 3), torch.rand(7, 3)]
padded_x = tk.pad({"input_ids": x}, padding=True)
print(padded_x)
breakpoint()