# I am lost on custom batch size definition

**URL:** <https://discuss.pytorch.org/t/i-am-lost-on-custom-batch-size-definition/179797>\
**Category:** data\
**Created:** [May 12, 2023, 7:07pm UTC](https://discuss.pytorch.org/t/i-am-lost-on-custom-batch-size-definition/179797 "2023-05-12T19:07:30Z")\
**Posts on this page:** 1\
**Showing post:** 6

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [May 14, 2023, 7:15pm UTC](https://discuss.pytorch.org/t/i-am-lost-on-custom-batch-size-definition/179797/6 "2023-05-14T19:15:14Z")

</div>

Sure, here is a simple example which uses pre-defined batch sizes in the `Dataset`:

```python
class MyDataset(Dataset):
    def __init__ (self):
        self.data = torch.randn(100, 1)
        self.target = torch.randint(0, 10, (100,))
        self.batch_sizes = [10, 20, 50, 15, 5]
        assert sum(self.batch_sizes) == self.data.size(0)
        
    def __len__ (self):
        return len(self.batch_sizes)
    
    def __getitem__ (self, index):
        batch_size = self.batch_sizes[index]
        offset = sum(self.batch_sizes[:index])
        print(f"loading {batch_size} samples at offset {offset}")
        x = []
        y = []
        for i in range(batch_size):
            x.append(self.data[offset+i])
            y.append(self.target[offset+i])
        x = torch.stack(x)
        y = torch.stack(y)
        return x, y

dataset = MyDataset()
loader = DataLoader(dataset, batch_size=1, shuffle=True)

for data, target in loader:
    # remove batch dimension added by the DataLoader
    data.squeeze_(0)
    target.squeeze_(0)
    print(data.shape, target.shape)
    
# loading 15 samples at offset 80
# torch.Size([15, 1]) torch.Size([15])
# loading 10 samples at offset 0
# torch.Size([10, 1]) torch.Size([10])
# loading 50 samples at offset 30
# torch.Size([50, 1]) torch.Size([50])
# loading 20 samples at offset 10
# torch.Size([20, 1]) torch.Size([20])
# loading 5 samples at offset 95
# torch.Size([5, 1]) torch.Size([5])

```

You could also take a look at the `BatchSampler` approach if it would fit your use case better as described in [this post](https://discuss.pytorch.org/t/how-to-use-batchsampler-with-getitem-dataset/78788/4).

---

_[View the full topic](https://discuss.pytorch.org/t/i-am-lost-on-custom-batch-size-definition/179797)._