sasan commited on
Commit
afa0e67
·
1 Parent(s): a3debde

chore: Remove commented out code in load_gpu_model function

Browse files
Files changed (1) hide show
  1. kitt/core/model.py +2 -2
kitt/core/model.py CHANGED
@@ -6,13 +6,13 @@ import xml.etree.ElementTree as ET
6
  from enum import Enum
7
  from typing import List
8
 
 
9
  from langchain.tools.base import StructuredTool
10
  from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
11
  from langchain_core.utils.function_calling import convert_to_openai_tool
12
  from loguru import logger
13
  from ollama import Client
14
  from pydantic import BaseModel
15
- import torch
16
 
17
  from kitt.skills import vehicle_status
18
  from kitt.skills.common import config
@@ -358,7 +358,7 @@ def load_gpu_model():
358
  device_map="auto",
359
  load_in_8bit=False,
360
  load_in_4bit=True,
361
- use_flash_attention_2=True,
362
  )
363
 
364
  return model, tokenizer
 
6
  from enum import Enum
7
  from typing import List
8
 
9
+ import torch
10
  from langchain.tools.base import StructuredTool
11
  from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
12
  from langchain_core.utils.function_calling import convert_to_openai_tool
13
  from loguru import logger
14
  from ollama import Client
15
  from pydantic import BaseModel
 
16
 
17
  from kitt.skills import vehicle_status
18
  from kitt.skills.common import config
 
358
  device_map="auto",
359
  load_in_8bit=False,
360
  load_in_4bit=True,
361
+ # use_flash_attention_2=True,
362
  )
363
 
364
  return model, tokenizer