def minibatch_gen()

in phasic_policy_gradient/minibatch_optimize.py [0:0]

• 12 lines of code
• 5 McCabe index (conditional complexity)

def minibatch_gen(data, *, batch_size=None, nminibatch=None, forever=False):
    assert (batch_size is None) != (
        nminibatch is None
    ), "only one of batch_size or nminibatch should be specified"
    ntrain = tu.batch_len(data)
    if nminibatch is None:
        nminibatch = max(ntrain // batch_size, 1)
    while True:
        for mbinds in th.chunk(th.randperm(ntrain), nminibatch):
            yield tree_map(to_th_device, tu.tree_slice(data, mbinds))
        if not forever:
            return