Hzfinfdu commited on
Commit
e3f0d7e
1 Parent(s): 45eceaf

Update quantization.py

Browse files
Files changed (1) hide show
  1. quantization.py +3 -7
quantization.py CHANGED
@@ -3,6 +3,9 @@ import torch
3
  import torch.nn as nn
4
  from torch.cuda.amp import custom_bwd, custom_fwd
5
  import math
 
 
 
6
 
7
 
8
  def find_layers(module, layers=[nn.Conv2d, nn.Linear], name=''):
@@ -16,13 +19,6 @@ def find_layers(module, layers=[nn.Conv2d, nn.Linear], name=''):
16
  return res
17
 
18
 
19
- try:
20
- import triton
21
- import triton.language as tl
22
- from .custom_autotune import *
23
- except:
24
- print('triton not installed. Run `pip install triton` to load quantized version of MOSS.')
25
-
26
  # code based https://github.com/fpgaminer/GPTQ-triton
27
  @autotune(
28
  configs=[
 
3
  import torch.nn as nn
4
  from torch.cuda.amp import custom_bwd, custom_fwd
5
  import math
6
+ import triton
7
+ import triton.language as tl
8
+ from .custom_autotune import *
9
 
10
 
11
  def find_layers(module, layers=[nn.Conv2d, nn.Linear], name=''):
 
19
  return res
20
 
21
 
 
 
 
 
 
 
 
22
  # code based https://github.com/fpgaminer/GPTQ-triton
23
  @autotune(
24
  configs=[