Make Python interface work for non-SM80 targets (#726)
* Make Python interface work for non-SM80 targets * Remove line in README
This commit is contained in:
@@ -34,6 +34,7 @@ import pycutlass
|
||||
from pycutlass import *
|
||||
from pycutlass.conv2d_operation import *
|
||||
from pycutlass.utils import reference_model
|
||||
from pycutlass.utils.device import device_cc
|
||||
import sys
|
||||
import torch.nn.functional as F
|
||||
|
||||
@@ -146,6 +147,11 @@ try:
|
||||
except:
|
||||
sys.exit(0)
|
||||
|
||||
cc = device_cc()
|
||||
if args.compute_capability != cc:
|
||||
raise Exception(("Parameter --compute-capability of {} "
|
||||
"does not match that of the device of {}.").format(args.compute_capability, cc))
|
||||
|
||||
pycutlass.get_memory_pool(init_pool_size=2**30, max_pool_size=2**32)
|
||||
|
||||
np.random.seed(0)
|
||||
|
||||
Reference in New Issue
Block a user