Searched refs:enable_nested_tensor (Results 1 – 2 of 2) sorted by relevance
| /aosp_15_r20/external/pytorch/test/ |
| H A D | test_transformers.py | 322 …def test_transformerencoder_fastpath(self, device, use_torchscript, enable_nested_tensor, use_auto… argument
338 enable_nested_tensor=enable_nested_tensor
413 …def test_transformerencoder_square_input(self, with_no_grad, training, enable_nested_tensor, devic… argument
421 enable_nested_tensor=enable_nested_tensor
458 def test_transformerencoder(self, batch_first, training, enable_nested_tensor, device): argument
488 def _test(batch_first, training, enable_nested_tensor): argument
496 encoder_layer, 1, enable_nested_tensor=enable_nested_tensor
560 …model = nn.TransformerEncoder(encoder_layer, 2, enable_nested_tensor=enable_nested_tensor).to(devi…
578 …model = nn.TransformerEncoder(encoder_layer, 6, enable_nested_tensor=enable_nested_tensor).to(devi…
600 enable_nested_tensor=enable_nested_tensor).to(device)
[all …]
|
| /aosp_15_r20/external/pytorch/torch/nn/modules/ |
| H A D | transformer.py | 336 enable_nested_tensor: bool = True,
345 self.enable_nested_tensor = enable_nested_tensor
347 self.use_nested_tensor = enable_nested_tensor
378 if enable_nested_tensor and why_not_sparsity_fast_path:
|