AncTe/ginka-generator
1
0
Fork 0
mirror of https://github.com/unanmed/ginka-generator.git synced 2026-05-14 04:41:12 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

use mps if available

Browse Source
This commit is contained in:
tocque 2026-02-06 23:21:59 +08:00
parent 9887abcd01
commit 2357363d76
1 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
ginka/train_vae.py
View File

@ -55,7 +55,11 @@ from shared.image import matrix_to_image_cv
BATCH_SIZE = 128 BATCH_SIZE = 128
LATENT_DIM = 48 LATENT_DIM = 48
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu") device = torch.device(
"cuda:1" if torch.cuda.is_available()
else "mps" if torch.mps.is_available()
else "cpu"
)
os.makedirs("result", exist_ok=True) os.makedirs("result", exist_ok=True)
os.makedirs("result/vae", exist_ok=True) os.makedirs("result/vae", exist_ok=True)
os.makedirs("result/ginka_vae_img", exist_ok=True) os.makedirs("result/ginka_vae_img", exist_ok=True)
@ -75,7 +79,7 @@ def parse_arguments():
return args return args
def train(): def train():
print(f"Using {'cuda' if torch.cuda.is_available() else 'cpu'} to train model.") print(f"Using {device.type} to train model.")
args = parse_arguments() args = parse_arguments()