text icon indicating copy to clipboard operation
text copied to clipboard

Fix docstring mistakes in bundler.py

Open yulonglin opened this issue 1 month ago • 1 comments

The output dimensions for one of the examples is wrong, since the batch size is 2, not 1. Further, there is no need to import RobertaClassificationHead for that example, as it is later accessed via torchtext.models.RobertaClassificationHead

How to reproduce the issue

Run the following:

>>> import torch, torchtext
>>> from torchtext.models import RobertaClassificationHead
>>> from torchtext.functional import to_tensor
>>> xlmr_large = torchtext.models.XLMR_LARGE_ENCODER
>>> classifier_head = torchtext.models.RobertaClassificationHead(num_classes=2, input_dim = 1024)
>>> model = xlmr_large.get_model(head=classifier_head)
>>> transform = xlmr_large.transform()
>>> input_batch = ["Hello world", "How are you!"]
>>> model_input = to_tensor(transform(input_batch), padding_value=1)
>>> output = model(model_input)
>>> output.shape

Expected result

torch.Size([2, 2])

Result given in docstring

torch.Size([1, 2])

yulonglin avatar May 16 '24 05:05 yulonglin