compile_complete_model.py

outputs = {}

#complete model input
inputs = Input(shape=x_train.shape[1:])
net = {'offload':inputs}

#stack all 3 device models together
for device in device_names:    
    net = device_models[device](net['offload'])
    outputs[device] = net[device+'_outputs']

#compile complete model
model = Model(inputs=inputs,outputs=outputs)
model.compile(loss=SparseCategoricalCrossentropy(from_logits=True),
              optimizer='nadam',
              metrics=['accuracy'])