Update model/openllama.py

model/openllama.py

@@ -11,7 +11,7 @@ import kornia as K
 import torch
 from torch.nn.utils import rnn
 from transformers import AutoConfig, AutoModelForCausalLM
-from accelerate import init_empty_weights, load_checkpoint_and_dispatch
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch, infer_auto_device_map
 
 CLASS_NAMES = ['bottle', 'cable', 'capsule', 'carpet', 'grid', 'hazelnut', 'leather', 'metal nut', 'pill', 'screw', 'tile', 'toothbrush', 'transistor', 'wood', 'zipper', 'object',
                'candle', 'cashew', 'chewinggum', 'fryum', 'macaroni', 'pcb', 'pipe fryum']