Home
last modified time | relevance | path

Searched refs:unsafe_set_bdim (Results 1 – 2 of 2) sorted by relevance

/external/pytorch/aten/src/ATen/functorch/
DBatchedTensorImpl.h97 void unsafe_set_bdim(int64_t bdim) { in unsafe_set_bdim() function
DLegacyBatchingRegistrations.cpp138 batched->unsafe_set_bdim(updated_batch_idx); in squeeze_dims__batching_rule()
180 batched->unsafe_set_bdim(new_bdim); in squeeze__batching_rule()
196 batched->unsafe_set_bdim(batched->bdim() + 1); in unsqueeze__batching_rule()