// logger.warning("Dumping to {}".format(""/Users/chaubold/Desktop/dump.h5""))
// vigra.impex.writeHDF5(reordered_feature_image[z,...], "data", "/Users/chaubold/Desktop/dump.h5")
result_slice = self._pytorch_net.forward([reordered_feature_image[z:z+1,...]])[0]
logger.info("Resulting slice {} has shape {}".format(z, result_slice.shape))
result[z, 0, ...] = result_slice
logger.info("Obtained a predicted block of shape {}".format(result.shape))
After Change
// vigra.impex.writeHDF5(reordered_feature_image[z,...], "data", "/Users/chaubold/Desktop/dump.h5")
// create batch of desired num slices. Multiple slices can be processed on multiple GPUs!
batch = [reordered_feature_image[zi:zi+1,...] for zi in range(z, min(z+self.BATCH_SIZE, reordered_feature_image.shape[0]))]
result_batch = self._pytorch_net.forward(batch)
logger.info("Resulting slices from {} to {} have shape {}".format(z, z+len(batch), result_batch[0].shape))
