no message

Author: sniklaus

correlation/correlation.py

@@ -246,11 +246,8 @@ def cupy_kernel(strFunction, objVariables):
 
         strTensor = objMatch.group(4)
         intSizes = objVariables[strTensor].size()
-        if torch.is_tensor(intSizes[intArg]):
-            strKernel = strKernel.replace(objectMatch.group(), str(intSizes[intArg].item()))
-        else:
-            strKernel = strKernel.replace(objectMatch.group(), str(intSizes[intArg]))
-        # end
+
+        strKernel = strKernel.replace(objMatch.group(), str(intSizes[intArg] if torch.is_tensor(intSizes[intArg]) == False else intSizes[intArg].item()))
 
     while True:
         objMatch = re.search('(VALUE_)([0-4])(\()([^\)]+)(\))', strKernel)
@@ -264,7 +261,7 @@ def cupy_kernel(strFunction, objVariables):
 
         strTensor = strArgs[0]
         intStrides = objVariables[strTensor].stride()
-        strIndex = [ '((' + strArgs[intArg + 1].replace('{', '(').replace('}', ')').strip() + ')*' + str(intStrides[intArg]) + ')' for intArg in range(intArgs) ]
+        strIndex = [ '((' + strArgs[intArg + 1].replace('{', '(').replace('}', ')').strip() + ')*' + str(intStrides[intArg] if torch.is_tensor(intStrides[intArg]) == False else intStrides[intArg].item()) + ')' for intArg in range(intArgs) ]
 
         strKernel = strKernel.replace(objMatch.group(0), strTensor + '[' + str.join('+', strIndex) + ']')
     # end