Skip to content

Commit

Permalink
cherry pick the tensor core fix (#536)
Browse files Browse the repository at this point in the history
Co-authored-by: Joel Schlosser <jbschlosser@fb.com>
  • Loading branch information
guotuofeng and jbschlosser authored Feb 18, 2022
1 parent cdfa44f commit 188c0c0
Showing 1 changed file with 1 addition and 2 deletions.
3 changes: 1 addition & 2 deletions tb_plugin/torch_tb_profiler/profiler/tensor_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,7 @@ def __contains__(cls, item):

class TC_OP_Allowlist(metaclass=TC_Allowlist_Meta):
# Refer to https://github.com/pytorch/pytorch/blob/69b2bf70f9c0e591ce5e566afa59e19618031ead/aten/src/ATen/autocast_mode.cpp#L290-L351 # noqa: E501
allowlist = ['aten::_convolution', 'aten::_convolution_nogroup',
'aten::conv1d', 'aten::conv2d', 'aten::conv3d', 'aten::conv_tbc',
allowlist = ['aten::_convolution', 'aten::conv1d', 'aten::conv2d', 'aten::conv3d', 'aten::conv_tbc',
'aten::conv_transpose1d', 'aten::conv_transpose2d', 'aten::conv_transpose3d',
'aten::convolution', 'aten::cudnn_convolution', 'aten::cudnn_convolution_transpose',
'aten::prelu', 'aten::addmm', 'aten::addmv', 'aten::addr',
Expand Down

0 comments on commit 188c0c0

Please sign in to comment.