check_errors (default: True): Checks if some neighbor pair was missed due to max_num_neighbors being too small. True is incompatible with CUDA-graphs.
static_shapes (default:False): Ensure tensornet uses statically shaped tensors regardless of the number of neighbors when true. Required for CUDA-graph compatibility.
The default values for these previous to this PR are tailored for inference (CUDA-graphs), but we have seen they can hurt performance and increase memory usage during training.
Thus, I am defaulting these to maximize training performance. Advanced users can tweak these to squeeze performance during inference.
cc @stefdoerr @guillemsimeon
guillemsimeon
commented
10 months ago
Adds two new options:
The default values for these previous to this PR are tailored for inference (CUDA-graphs), but we have seen they can hurt performance and increase memory usage during training. Thus, I am defaulting these to maximize training performance. Advanced users can tweak these to squeeze performance during inference. cc @stefdoerr @guillemsimeon