Invalid argument: Assign requires shapes of both tensors to match. lhs shape= [10,128] rhs shape= [32,128]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): NO
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04, aarch64, Nvidia Jetson Nano
• Mobile device (e.g., Pixel 4, Samsung Galaxy 10) if the issue happens on mobile device: n.a.
• TensorFlow installed from (source or binary): Binary
• TensorFlow version (use command below): bcudart.so.10.0, unknown 1.15.0 - 1.15.0+nv19.11.tf1
• Python version: Python 3.6.9
• Bazel version (if compiling from source): n.a.
• GCC/Compiler version (if compiling from source): n.a.
• CUDA/cuDNN version: 10.0 / 7.5
• GPU model and memory: maxwell, 4 GB

Please provide the entire URL of the model you are using?
https://github.com/tensorflow/magenta/tree/master/magenta/models/image_stylization

Describe the current behavior
Crashes

Describe the expected behavior
Generating graphics :-)

Code to reproduce the issue
image_stylization_transform
--num_styles=10
--checkpoint=/models/imstyle/multistyle-pastiche-generator-varied.ckpt
--input_image=/pics/Sarah.jpg
--which_styles="[$i]"
--output_dir=/pics
--output_basename="varied_"

Other info / logs
Only occurs with the 'varied' model. The 'monet' model works fine.

2020-03-03 18:53:09.297669: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcudart.so.10.0
2020-03-03 18:53:22.827005: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcuda.so.1
2020-03-03 18:53:22.839411: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:950] ARM64 does not support NUMA - returning NUMA node zero
2020-03-03 18:53:22.839564: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1618] Found device 0 with properties:
name: NVIDIA Tegra X1 major: 5 minor: 3 memoryClockRate(GHz): 0.9216
pciBusID: 0000:00:00.0
2020-03-03 18:53:22.839628: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcudart.so.10.0
2020-03-03 18:53:22.843182: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcublas.so.10.0
2020-03-03 18:53:22.845692: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcufft.so.10.0
2020-03-03 18:53:22.846743: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcurand.so.10.0
2020-03-03 18:53:22.850327: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcusolver.so.10.0
2020-03-03 18:53:22.853101: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcusparse.so.10.0
2020-03-03 18:53:22.862269: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcudnn.so.7
2020-03-03 18:53:22.862856: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:950] ARM64 does not support NUMA - returning NUMA node zero
2020-03-03 18:53:22.863368: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:950] ARM64 does not support NUMA - returning NUMA node zero
2020-03-03 18:53:22.863474: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1746] Adding visible gpu devices: 0
2020-03-03 18:53:22.885156: W tensorflow/core/platform/profile_utils/cpu_utils.cc:98] Failed to find bogomips in /proc/cpuinfo; cannot determine CPU frequency
2020-03-03 18:53:22.885953: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x33177100 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2020-03-03 18:53:22.886010: I tensorflow/compiler/xla/service/service.cc:176] StreamExecutor device (0): Host, Default Version
2020-03-03 18:53:22.957976: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:950] ARM64 does not support NUMA - returning NUMA node zero
2020-03-03 18:53:22.958290: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x330f65b0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2020-03-03 18:53:22.958340: I tensorflow/compiler/xla/service/service.cc:176] StreamExecutor device (0): NVIDIA Tegra X1, Compute Capability 5.3
2020-03-03 18:53:22.959034: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:950] ARM64 does not support NUMA - returning NUMA node zero
2020-03-03 18:53:22.959154: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1618] Found device 0 with properties:
name: NVIDIA Tegra X1 major: 5 minor: 3 memoryClockRate(GHz): 0.9216
pciBusID: 0000:00:00.0
2020-03-03 18:53:22.959227: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcudart.so.10.0
2020-03-03 18:53:22.959353: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcublas.so.10.0
2020-03-03 18:53:22.959434: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcufft.so.10.0
2020-03-03 18:53:22.959506: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcurand.so.10.0
2020-03-03 18:53:22.959576: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcusolver.so.10.0
2020-03-03 18:53:22.959643: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcusparse.so.10.0
2020-03-03 18:53:22.959712: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcudnn.so.7
2020-03-03 18:53:22.960011: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:950] ARM64 does not support NUMA - returning NUMA node zero
2020-03-03 18:53:22.960335: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:950] ARM64 does not support NUMA - returning NUMA node zero
2020-03-03 18:53:22.960416: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1746] Adding visible gpu devices: 0
2020-03-03 18:53:22.960493: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcudart.so.10.0
2020-03-03 18:53:24.233468: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1159] Device interconnect StreamExecutor with strength 1 edge matrix:
2020-03-03 18:53:24.233555: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1165] 0
2020-03-03 18:53:24.233588: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1178] 0: N
2020-03-03 18:53:24.234153: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:950] ARM64 does not support NUMA - returning NUMA node zero
2020-03-03 18:53:24.234584: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:950] ARM64 does not support NUMA - returning NUMA node zero
2020-03-03 18:53:24.234801: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1304] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 1757 MB memory) -> physical GPU (device: 0, name: NVIDIA Tegra X1, pci bus id: 0000:00:00.0, compute capability: 5.3)
WARNING:tensorflow:From /home/lex/python-envs/m/lib/python3.6/site-packages/tensorflow_core/contrib/layers/python/layers/layers.py:1060: Layer.apply (from tensorflow.python.keras.engine.base_layer) is deprecated and will be removed in a future version.
Instructions for updating:
Please use layer.__call__ method instead.
W0303 18:53:24.298795 547830423568 deprecation.py:323] From /home/lex/python-envs/m/lib/python3.6/site-packages/tensorflow_core/contrib/layers/python/layers/layers.py:1060: Layer.apply (from tensorflow.python.keras.engine.base_layer) is deprecated and will be removed in a future version.
Instructions for updating:
Please use layer.__call__ method instead.
INFO:tensorflow:Restoring parameters from /models/imstyle/multistyle-pastiche-generator-varied.ckpt
I0303 18:53:27.952254 547830423568 saver.py:1284] Restoring parameters from /models/imstyle/multistyle-pastiche-generator-varied.ckpt
Traceback (most recent call last):
File "/home/lex/python-envs/m/lib/python3.6/site-packages/tensorflow_core/python/client/session.py", line 1365, in _do_call
return fn(*args)
File "/home/lex/python-envs/m/lib/python3.6/site-packages/tensorflow_core/python/client/session.py", line 1350, in _run_fn
target_list, run_metadata)
File "/home/lex/python-envs/m/lib/python3.6/site-packages/tensorflow_core/python/client/session.py", line 1443, in _call_tf_sessionrun
run_metadata)
tensorflow.python.framework.errors_impl.InvalidArgumentError: 2 root error(s) found.
(0) Invalid argument: Assign requires shapes of both tensors to match. lhs shape= [10,128] rhs shape= [32,128]
[[{{node save/Assign_46}}]]
(1) Invalid argument: Assign requires shapes of both tensors to match. lhs shape= [10,128] rhs shape= [32,128]
[[{{node save/Assign_46}}]]
[[save/RestoreV2/_2]]
0 successful operations.
0 derived errors ignored.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/home/lex/python-envs/m/lib/python3.6/site-packages/tensorflow_core/python/training/saver.py", line 1290, in restore
{self.saver_def.filename_tensor_name: save_path})
File "/home/lex/python-envs/m/lib/python3.6/site-packages/tensorflow_core/python/client/session.py", line 956, in run
run_metadata_ptr)
File "/home/lex/python-envs/m/lib/python3.6/site-packages/tensorflow_core/python/client/session.py", line 1180, in _run
feed_dict_tensor, options, run_metadata)
File "/home/lex/python-envs/m/lib/python3.6/site-packages/tensorflow_core/python/client/session.py", line 1359, in _do_run
run_metadata)
File "/home/lex/python-envs/m/lib/python3.6/site-packages/tensorflow_core/python/client/session.py", line 1384, in _do_call
raise type(e)(node_def, op, message)
tensorflow.python.framework.errors_impl.InvalidArgumentError: 2 root error(s) found.
(0) Invalid argument: Assign requires shapes of both tensors to match. lhs shape= [10,128] rhs shape= [32,128]
[[node save/Assign_46 (defined at /lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py:1748) ]]
(1) Invalid argument: Assign requires shapes of both tensors to match. lhs shape= [10,128] rhs shape= [32,128]
[[node save/Assign_46 (defined at /lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py:1748) ]]
[[save/RestoreV2/_2]]
0 successful operations.
0 derived errors ignored.

Original stack trace for 'save/Assign_46':
File "/bin/image_stylization_transform", line 8, in
sys.exit(console_entry_point())
File "/lib/python3.6/site-packages/magenta/models/image_stylization/image_stylization_transform.py", line 147, in console_entry_point
tf.app.run(main)
File "/lib/python3.6/site-packages/tensorflow_core/python/platform/app.py", line 40, in run
_run(main=main, argv=argv, flags_parser=_parse_flags_tolerate_undef)
File "/lib/python3.6/site-packages/absl/app.py", line 299, in run
_run_main(main, args)
File "/lib/python3.6/site-packages/absl/app.py", line 250, in _run_main
sys.exit(main(argv))
File "/lib/python3.6/site-packages/magenta/models/image_stylization/image_stylization_transform.py", line 138, in main
_multiple_images(image, which_styles, output_dir)
File "/lib/python3.6/site-packages/magenta/models/image_stylization/image_stylization_transform.py", line 95, in _multiple_images
_load_checkpoint(sess, FLAGS.checkpoint)
File "/lib/python3.6/site-packages/magenta/models/image_stylization/image_stylization_transform.py", line 51, in _load_checkpoint
model_saver = tf.train.Saver(tf.global_variables())
File "/lib/python3.6/site-packages/tensorflow_core/python/training/saver.py", line 828, in init
self.build()
File "/lib/python3.6/site-packages/tensorflow_core/python/training/saver.py", line 840, in build
self._build(self._filename, build_save=True, build_restore=True)
File "/lib/python3.6/site-packages/tensorflow_core/python/training/saver.py", line 878, in _build
build_restore=build_restore)
File "/lib/python3.6/site-packages/tensorflow_core/python/training/saver.py", line 508, in _build_internal
restore_sequentially, reshape)
File "/lib/python3.6/site-packages/tensorflow_core/python/training/saver.py", line 350, in _AddRestoreOps
assign_ops.append(saveable.restore(saveable_tensors, shapes))
File "/lib/python3.6/site-packages/tensorflow_core/python/training/saving/saveable_object_util.py", line 73, in restore
self.op.get_shape().is_fully_defined())
File "/lib/python3.6/site-packages/tensorflow_core/python/ops/state_ops.py", line 227, in assign
validate_shape=validate_shape)
File "/lib/python3.6/site-packages/tensorflow_core/python/ops/gen_state_ops.py", line 66, in assign
use_locking=use_locking, name=name)
File "/lib/python3.6/site-packages/tensorflow_core/python/framework/op_def_library.py", line 794, in _apply_op_helper
op_def=op_def)
File "/lib/python3.6/site-packages/tensorflow_core/python/util/deprecation.py", line 507, in new_func
return func(*args, **kwargs)
File "/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py", line 3357, in create_op
attrs, op_def, compute_device)
File "/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py", line 3426, in _create_op_internal
op_def=op_def)
File "/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py", line 1748, in init
self._traceback = tf_stack.extract_stack()

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/home/lex/python-envs/m/bin/image_stylization_transform", line 8, in
sys.exit(console_entry_point())
File "/home/lex/python-envs/m/lib/python3.6/site-packages/magenta/models/image_stylization/image_stylization_transform.py", line 147, in console_entry_point
tf.app.run(main)
File "/home/lex/python-envs/m/lib/python3.6/site-packages/tensorflow_core/python/platform/app.py", line 40, in run
_run(main=main, argv=argv, flags_parser=_parse_flags_tolerate_undef)
File "/home/lex/python-envs/m/lib/python3.6/site-packages/absl/app.py", line 299, in run
_run_main(main, args)
File "/home/lex/python-envs/m/lib/python3.6/site-packages/absl/app.py", line 250, in _run_main
sys.exit(main(argv))
File "/home/lex/python-envs/m/lib/python3.6/site-packages/magenta/models/image_stylization/image_stylization_transform.py", line 138, in main
_multiple_images(image, which_styles, output_dir)
File "/home/lex/python-envs/m/lib/python3.6/site-packages/magenta/models/image_stylization/image_stylization_transform.py", line 95, in _multiple_images
_load_checkpoint(sess, FLAGS.checkpoint)
File "/home/lex/python-envs/m/lib/python3.6/site-packages/magenta/models/image_stylization/image_stylization_transform.py", line 56, in _load_checkpoint
model_saver.restore(sess, checkpoint)
File "/home/lex/python-envs/m/lib/python3.6/site-packages/tensorflow_core/python/training/saver.py", line 1326, in restore
err, "a mismatch between the current graph and the graph")
tensorflow.python.framework.errors_impl.InvalidArgumentError: Restoring from checkpoint failed. This is most likely due to a mismatch between the current graph and the graph from the checkpoint. Please ensure that you have not altered the graph expected based on the checkpoint. Original error:

2 root error(s) found.
(0) Invalid argument: Assign requires shapes of both tensors to match. lhs shape= [10,128] rhs shape= [32,128]
[[node save/Assign_46 (defined at /lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py:1748) ]]
(1) Invalid argument: Assign requires shapes of both tensors to match. lhs shape= [10,128] rhs shape= [32,128]
[[node save/Assign_46 (defined at /lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py:1748) ]]
[[save/RestoreV2/_2]]
0 successful operations.
0 derived errors ignored.

Original stack trace for 'save/Assign_46':
File "/bin/image_stylization_transform", line 8, in
sys.exit(console_entry_point())
File "/lib/python3.6/site-packages/magenta/models/image_stylization/image_stylization_transform.py", line 147, in console_entry_point
tf.app.run(main)
File "/lib/python3.6/site-packages/tensorflow_core/python/platform/app.py", line 40, in run
_run(main=main, argv=argv, flags_parser=_parse_flags_tolerate_undef)
File "/lib/python3.6/site-packages/absl/app.py", line 299, in run
_run_main(main, args)
File "/lib/python3.6/site-packages/absl/app.py", line 250, in _run_main
sys.exit(main(argv))
File "/lib/python3.6/site-packages/magenta/models/image_stylization/image_stylization_transform.py", line 138, in main
_multiple_images(image, which_styles, output_dir)
File "/lib/python3.6/site-packages/magenta/models/image_stylization/image_stylization_transform.py", line 95, in _multiple_images
_load_checkpoint(sess, FLAGS.checkpoint)
File "/lib/python3.6/site-packages/magenta/models/image_stylization/image_stylization_transform.py", line 51, in _load_checkpoint
model_saver = tf.train.Saver(tf.global_variables())
File "/lib/python3.6/site-packages/tensorflow_core/python/training/saver.py", line 828, in init
self.build()
File "/lib/python3.6/site-packages/tensorflow_core/python/training/saver.py", line 840, in build
self._build(self._filename, build_save=True, build_restore=True)
File "/lib/python3.6/site-packages/tensorflow_core/python/training/saver.py", line 878, in _build
build_restore=build_restore)
File "/lib/python3.6/site-packages/tensorflow_core/python/training/saver.py", line 508, in _build_internal
restore_sequentially, reshape)
File "/lib/python3.6/site-packages/tensorflow_core/python/training/saver.py", line 350, in _AddRestoreOps
assign_ops.append(saveable.restore(saveable_tensors, shapes))
File "/lib/python3.6/site-packages/tensorflow_core/python/training/saving/saveable_object_util.py", line 73, in restore
self.op.get_shape().is_fully_defined())
File "/lib/python3.6/site-packages/tensorflow_core/python/ops/state_ops.py", line 227, in assign
validate_shape=validate_shape)
File "/lib/python3.6/site-packages/tensorflow_core/python/ops/gen_state_ops.py", line 66, in assign
use_locking=use_locking, name=name)
File "/lib/python3.6/site-packages/tensorflow_core/python/framework/op_def_library.py", line 794, in _apply_op_helper
op_def=op_def)
File "/lib/python3.6/site-packages/tensorflow_core/python/util/deprecation.py", line 507, in new_func
return func(*args, **kwargs)
File "/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py", line 3357, in create_op
attrs, op_def, compute_device)
File "/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py", line 3426, in _create_op_internal
op_def=op_def)
File "/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py", line 1748, in init
self._traceback = tf_stack.extract_stack()