From 3695fd996c94557a72e78a1f22055e82afbe0955 Mon Sep 17 00:00:00 2001
From: Gitea <gitea@fake.local>
Date: Fri, 21 Oct 2022 16:40:38 +0800
Subject: [PATCH] added

---
 models/ai_model_manage.go        | 2 +-
 routers/repo/ai_model_convert.go | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/models/ai_model_manage.go b/models/ai_model_manage.go
index a88da8fe5..d9adda2dc 100644
--- a/models/ai_model_manage.go
+++ b/models/ai_model_manage.go
@@ -88,7 +88,7 @@ type AiModelQueryOptions struct {
 }
 
 func (a *AiModelConvert) IsGpuTrainTask() bool {
-	if a.SrcEngine == 0 || a.SrcEngine == 1 {
+	if a.SrcEngine == 0 || a.SrcEngine == 1 || a.SrcEngine == 4 || a.SrcEngine == 6 {
 		return true
 	}
 	return false
diff --git a/routers/repo/ai_model_convert.go b/routers/repo/ai_model_convert.go
index bd6a01072..13b02f9d0 100644
--- a/routers/repo/ai_model_convert.go
+++ b/routers/repo/ai_model_convert.go
@@ -18,7 +18,7 @@ import (
 	"code.gitea.io/gitea/modules/git"
 	"code.gitea.io/gitea/modules/log"
 	"code.gitea.io/gitea/modules/modelarts"
-	"code.gitea.io/gitea/modules/setting"
+	"code.gitea.io/gitea/modules/set ting"
 	"code.gitea.io/gitea/modules/storage"
 	"code.gitea.io/gitea/modules/timeutil"
 	uuid "github.com/satori/go.uuid"