Constructs a version of 'fn' that applies to smaller batches.
(fn, chunk)
| 24 | |
| 25 | |
| 26 | def batchify(fn, chunk): |
| 27 | """Constructs a version of 'fn' that applies to smaller batches. |
| 28 | """ |
| 29 | if chunk is None: |
| 30 | return fn |
| 31 | def ret(inputs): |
| 32 | return torch.cat([fn(inputs[i:i+chunk]) for i in range(0, inputs.shape[0], chunk)], 0) |
| 33 | return ret |
| 34 | |
| 35 | |
| 36 | def run_network(inputs, viewdirs, fn, embed_fn, embeddirs_fn, netchunk=1024*64): |