fa6bd25500ef229d03102aa49d71b969d64f1f09,app/websocket/app/main/persistence/utils.py,,snapshot,#Any#,8

Before Change

    // Add new objects from database
    new_keys = current_keys - last_snapshot_keys
    objects = [
        WorkerObject(worker_id=worker.id, object=worker._objects[key], id=key)
        for key in new_keys
    ]

After Change

        obj = worker.get_obj(key)

        // If obj is a Plan we save in the database as a model
        if isinstance(obj, Plan):
            models.append(TorchModel(id=key, model=sy.serde.serialize(obj)))
        // If obj is Jit we ignore it
        elif isinstance(obj, torch.jit.ScriptModule):
            continue
        // If obj is a parameter we wrap the data and store it in the database
        // as an object
        elif isinstance(obj, torch.nn.Parameter):
            obj = obj.data.wrap()
            objects.append(WorkerObject(worker_id=worker.id, object=obj, id=key))
        // Otherwise we just store the object in the database
        else:
            objects.append(WorkerObject(worker_id=worker.id, object=obj, id=key))

    db.session.add_all(objects)
    db.session.add_all(models)
    db.session.commit()
    last_snapshot_keys = current_keys

Instances