Not able to fine-tune on fashion Try-on

[ManuelCecere]

Hello there, thanks for all the fantastic work done for AnyDoor!
I'm trying to fine-tune the model for fashion try-on application, but when I try to unfreeze just the ResBlock or the SpatialTransformer module and not both together a RuntimeError One of the differentiated Tensors does not require grad is raised.
That's troublesome, even though now I'm working with both unfreeze, because I was planning of targeting just the convolutional and linear layers, using LoRA but doing so raises the error.
Any tip would be greatly appreciated, thanks!

The code for reference, in the run_train_anydoor script:

.... the code above is unaltered
model = create_model("./configs/anydoor.yaml").cpu()
model.load_state_dict(load_state_dict(resume_path, location="cpu"))
model.learning_rate = learning_rate
model.sd_locked = sd_locked
model.only_mid_control = only_mid_control

for name, param in model.named_parameters():
param.requires_grad = False

for block in model.model.diffusion_model.output_blocks:
for layer in block:
if isinstance(layer, SpatialTransformer):
for name, param in layer.named_parameters():
param.requires_grad = True

dataloader = DataLoader(dataset, num_workers=8, batch_size=batch_size, shuffle=True)
logger = ImageLogger(batch_frequency=logger_freq)
trainer = pl.Trainer(
gpus=n_gpus,
strategy="ddp",
precision=16,
accelerator="gpu",
callbacks=[logger],
progress_bar_refresh_rate=1,
accumulate_grad_batches=accumulate_grad_batches,
)

trainer.fit(model, dataloader)