你抛出了StopIteration异常，但是忘记了重置下一个epoch的索引！因此，它会在一个epoch后自动停止。
我已经将您的代码片段扩展为一个工作代码示例（没有错别字），它应该以您预期的方式工作。

import torch
import numpy as np
from torch.utils.data import Sampler
from torch.utils.data import DataLoader, TensorDataset


class VariableBatchSampler(Sampler):
    def __init__(self, dataset_len: int, batch_sizes: list):
        self.dataset_len = dataset_len
        self.batch_sizes = batch_sizes
        self.batch_idx = 0
        self.start_idx = 0
        self.end_idx = self.batch_sizes[self.batch_idx]

    def __iter__(self):
        return self

    def __next__(self):
        if self.start_idx >= self.dataset_len:
            self.batch_idx = 0
            self.start_idx = 0
            self.end_idx = self.batch_sizes[self.batch_idx]
            raise StopIteration

        batch_indices = list(range(self.start_idx, self.end_idx))
        self.start_idx = self.end_idx
        self.batch_idx += 1

        try:
            self.end_idx += self.batch_sizes[self.batch_idx]
        except IndexError:
            self.end_idx = self.dataset_len

        return batch_indices

x_train = torch.randn(23)
y_train = torch.randint(0, 2, (23,))

batch_sizes = [4, 10, 7, 2]
train_dataset = TensorDataset(x_train, y_train)
sampler = VariableBatchSampler(dataset_len=len(x_train), batch_sizes=batch_sizes)
dataloader_train = DataLoader(train_dataset, sampler=sampler)

max_epoch = 4
for epoch in np.arange(1, max_epoch):
    print("Epoch: ", epoch)
    for x_batch, y_batch in dataloader_train:
         print(x_batch.shape)

字符串
这将产生：

Epoch: 1
torch.Size([1, 4])
torch.Size([1, 10])
torch.Size([1, 7])
torch.Size([1, 2])
Epoch: 2
torch.Size([1, 4])
torch.Size([1, 10])
torch.Size([1, 7])
torch.Size([1, 2])
Epoch: 3
torch.Size([1, 4])
torch.Size([1, 10])
torch.Size([1, 7])
torch.Size([1, 2])

型