Searched defs:need_weights (Results 1 – 6 of 6) sorted by relevance
| /external/pytorch/test/ |
| D | test_native_mha.py | 113 …self, device, dtype, mode, use_nt, need_weights, average_attn_weights, use_padding=False, pad_all=… argument
278 … need_weights, average_attn_weights, use_padding, pad_all, fused): argument
|
| D | test_jit.py | 14923 need_weights=True, # type: bool argument
|
| /external/pytorch/torch/csrc/api/src/nn/modules/ |
| D | activation.cpp | 445 bool need_weights, in forward()
|
| /external/pytorch/aten/src/ATen/native/transformers/ |
| D | attention.cpp | 272 bool need_weights, in native_multi_head_attention_cpu()
|
| /external/pytorch/aten/src/ATen/native/transformers/cuda/ |
| D | attention.cu | 490 bool need_weights, in native_multi_head_attention_cuda()
|
| /external/pytorch/test/inductor/ |
| D | test_cpu_repro.py | 302 need_weights, argument
|