text
text copied to clipboard
Fix docstring mistakes in bundler.py
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])