ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-30 19:21:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added inbuild plugins and made one for image referenced. WIP

Browse Source
This commit is contained in:
Jaret Burkett
2023-08-10 16:02:44 -06:00
parent df48f0a843
commit 1a7e346b41
12 changed files with 338 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
jobs/process/TrainSDRescaleProcess.py
View File

@@ -210,7 +210,7 @@ class TrainSDRescaleProcess(BaseSDTrainProcess):
flush()
# end hook_before_train_loop
def hook_train_loop(self):
def hook_train_loop(self, batch):
dtype = get_torch_dtype(self.train_config.dtype)
loss_function = torch.nn.MSELoss()