# Buggy code
def _iterbatches_from_shards(
    self, shard_indices, batch_size=None, deterministic=False, pad_batches=False
):
    """Get an object that iterates over batches from a restricted set of shards."""

    def iterate(dataset, batch_size):
        num_shards = len(shard_indices)
        if not deterministic:
            shard_perm = np.random.permutation(num_shards)
        else:
            shard_perm = np.arange(num_shards)

        # (ytz): Depending on the application, thread-based pools may be faster
        # than process based pools, since process based pools need to pickle/serialize
        # objects as an extra overhead. Also, as hideously as un-thread safe this looks,
        # we're actually protected by the GIL.
        pool = Pool(1)  # mp.dummy aliases ThreadPool to Pool
        next_shard = pool.apply_async(
            dataset.get_shard, (shard_indices[shard_perm[0]],)
        )

        total_yield = 0

        if batch_size is None:
            num_global_batches = num_shards
        else:
            num_global_batches = math.ceil(dataset.get_shape()[0][0] / batch_size)

        cur_global_batch = 0
        cur_shard = 0
        carry = None

        while cur_global_batch < num_global_batches:
            X, y, w, ids = next_shard.get()
            if cur_shard < num_shards - 1:
                next_shard = pool.apply_async(
                    dataset.get_shard, (shard_indices[shard_perm[cur_shard + 1]],)
                )
            else:
                pool.close()

            if carry is not None:
                X = np.concatenate([carry[0], X], axis=0)
                if y is not None:
                    y = np.concatenate([carry[1], y], axis=0)
                if w is not None:
                    w = np.concatenate([carry[2], w], axis=0)
                ids = np.concatenate([carry[3], ids], axis=0)
                carry = None

            n_shard_samples = X.shape[0]
            cur_local_batch = 0
            if batch_size is None:
                shard_batch_size = n_shard_samples
            else:
                shard_batch_size = batch_size

            if n_shard_samples == 0:
                cur_shard += 1
                if batch_size is None:
                    cur_global_batch += 1
                continue

            num_local_batches = math.ceil(n_shard_samples / shard_batch_size)
            if not deterministic:
                sample_perm = np.random.permutation(n_shard_samples)
            else:
                sample_perm = np.arange(n_shard_samples)

            while cur_local_batch < num_local_batches:
                start = cur_local_batch * shard_batch_size
                end = min(n_shard_samples, (cur_local_batch + 1) * shard_batch_size)

                indices = range(start, end)
                perm_indices = sample_perm[indices]
                X_b = X[perm_indices]

                if y is not None:
                    y_b = y[perm_indices]
                else:
                    y_b = None

                if w is not None:
                    w_b = w[perm_indices]
                else:
                    w_b = None

                ids_b = ids[perm_indices]

                assert len(X_b) <= shard_batch_size
                if len(X_b) < shard_batch_size and cur_shard != num_shards - 1:
                    assert carry is None
                    carry = [X_b, y_b, w_b, ids_b]
                else:
                    # (ytz): this skips everything except possibly the last shard
                    if pad_batches:
                        (X_b, y_b, w_b, ids_b) = pad_batch(
                            shard_batch_size, X_b, y_b, w_b, ids_b
                        )

                    yield X_b, y_b, w_b, ids_b
                    cur_global_batch += 1
                cur_local_batch += 1
            cur_shard += 1

    return iterate(self, batch_size)


# Traceback (AttributeError)
Traceback (most recent call last):
File "test_fit_on_batch.py", line 23, in <module>
model.fit_on_batch(X, y, w)
File "/Users/bharath/Code/deepchem/deepchem/models/keras_model.py", line 428, in fit_on_batch
if not self.built:
AttributeError: 'GraphConvModel' object has no attribute 'built'

# Fixed code
def _iterbatches_from_shards(
    self,
    shard_indices,
    batch_size=None,
    epochs=1,
    deterministic=False,
    pad_batches=False,
):
    """Get an object that iterates over batches from a restricted set of shards."""

    def iterate(dataset, batch_size, epochs):
        num_shards = len(shard_indices)
        if deterministic:
            shard_perm = np.arange(num_shards)

        # (ytz): Depending on the application, thread-based pools may be faster
        # than process based pools, since process based pools need to pickle/serialize
        # objects as an extra overhead. Also, as hideously as un-thread safe this looks,
        # we're actually protected by the GIL.
        pool = Pool(1)  # mp.dummy aliases ThreadPool to Pool

        if batch_size is None:
            num_global_batches = num_shards
        else:
            num_global_batches = math.ceil(dataset.get_shape()[0][0] / batch_size)

        for epoch in range(epochs):
            if not deterministic:
                shard_perm = np.random.permutation(num_shards)
            next_shard = pool.apply_async(
                dataset.get_shard, (shard_indices[shard_perm[0]],)
            )
            cur_global_batch = 0
            cur_shard = 0
            carry = None

            while cur_global_batch < num_global_batches:
                X, y, w, ids = next_shard.get()
                if cur_shard < num_shards - 1:
                    next_shard = pool.apply_async(
                        dataset.get_shard, (shard_indices[shard_perm[cur_shard + 1]],)
                    )
                elif epoch == epochs - 1:
                    pool.close()

                if carry is not None:
                    X = np.concatenate([carry[0], X], axis=0)
                    if y is not None:
                        y = np.concatenate([carry[1], y], axis=0)
                    if w is not None:
                        w = np.concatenate([carry[2], w], axis=0)
                    ids = np.concatenate([carry[3], ids], axis=0)
                    carry = None

                n_shard_samples = X.shape[0]
                cur_local_batch = 0
                if batch_size is None:
                    shard_batch_size = n_shard_samples
                else:
                    shard_batch_size = batch_size

                if n_shard_samples == 0:
                    cur_shard += 1
                    if batch_size is None:
                        cur_global_batch += 1
                    continue

                num_local_batches = math.ceil(n_shard_samples / shard_batch_size)
                if not deterministic:
                    sample_perm = np.random.permutation(n_shard_samples)
                else:
                    sample_perm = np.arange(n_shard_samples)

                while cur_local_batch < num_local_batches:
                    start = cur_local_batch * shard_batch_size
                    end = min(n_shard_samples, (cur_local_batch + 1) * shard_batch_size)

                    indices = range(start, end)
                    perm_indices = sample_perm[indices]
                    X_b = X[perm_indices]

                    if y is not None:
                        y_b = y[perm_indices]
                    else:
                        y_b = None

                    if w is not None:
                        w_b = w[perm_indices]
                    else:
                        w_b = None

                    ids_b = ids[perm_indices]

                    assert len(X_b) <= shard_batch_size
                    if len(X_b) < shard_batch_size and cur_shard != num_shards - 1:
                        assert carry is None
                        carry = [X_b, y_b, w_b, ids_b]
                    else:
                        # (ytz): this skips everything except possibly the last shard
                        if pad_batches:
                            (X_b, y_b, w_b, ids_b) = pad_batch(
                                shard_batch_size, X_b, y_b, w_b, ids_b
                            )

                        yield X_b, y_b, w_b, ids_b
                        cur_global_batch += 1
                    cur_local_batch += 1
                cur_shard += 1

    return iterate(self, batch_size, epochs)