hongshaorou/dreamer4
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

tweak bc trainer

Browse Source
This commit is contained in:
lucidrains 2025-10-21 10:54:47 -07:00
parent 2fc3b17149
commit e280592510
2 changed files with 9 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
dreamer4/trainers.py
View File

@ -195,7 +195,14 @@ class BehaviorCloneTrainer(Module):
for _ in range(self.num_train_steps): for _ in range(self.num_train_steps):
batch_data = next(iter_train_dl) batch_data = next(iter_train_dl)
loss = self.model(**batch_data) # just assume raw video dynamics training if batch_data is a tensor
# else kwargs for video, actions, rewards
if is_tensor(batch_data):
loss = self.model(batch_data)
else:
loss = self.model(**batch_data)
self.accelerator.backward(loss) self.accelerator.backward(loss)
if exists(self.max_grad_norm): if exists(self.max_grad_norm):

2
pyproject.toml
View File

@ -1,6 +1,6 @@
[project] [project]
name = "dreamer4" name = "dreamer4"
version = "0.0.57" version = "0.0.58"
description = "Dreamer 4" description = "Dreamer 4"
authors = [ authors = [
{ name = "Phil Wang", email = "lucidrains@gmail.com" } { name = "Phil Wang", email = "lucidrains@gmail.com" }