我的程序在两个 GPU 上都能正常工作,但在到达 torch.bmm 时会引发此错误
energy = torch.bmm(proj_query,proj_key) # transpose check
RuntimeError: cublas runtime error : the GPU program failed to execute at /pytorch/aten/src/THC/THCBlas.cu:450
我该如何解决这个问题?
我的程序在两个 GPU 上都能正常工作,但在到达 torch.bmm 时会引发此错误
energy = torch.bmm(proj_query,proj_key) # transpose check
RuntimeError: cublas runtime error : the GPU program failed to execute at /pytorch/aten/src/THC/THCBlas.cu:450
我该如何解决这个问题?