RuntimeError when running baselines/imagenet/sngp.py about uncertainty-baselines

[dustinvtran]

@jereliu

Hi @batzner! Were you able to resolve this?

[batzner]

Unfortunately, I was not able to resolve this. I saw in GitHub issues in other repos that people who got the same error message were able to resolve it by specifying synchronization=tf.VariableSynchronization.ON_READ in the self.add_weight(...) call (see below).

But I am not sure whether the on-read-synchronization is the correct behavior in this case.

self.precision_matrix = (
    self.add_weight(
        name='gp_precision_matrix',
        shape=(gp_feature_dim, gp_feature_dim),
        dtype=self.dtype,
        initializer=tf.keras.initializers.Identity(self.ridge_penalty),
        trainable=False,
        aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA))

(from https://github.com/google/edward2/blob/807bd74d93c607a5a4030c4ef7debecf89f8b6ab/edward2/tensorflow/layers/random_feature.py#L337)

[znado]

Hey! Sorry for the delay.

What was the command you were using to run the code, and what was the GPU/version of CUDA/version of TF? Also, have you been able to reproduce this at the current HEAD?

[gpleiss]

I get the exact same error. Fresh install (nightly version of tf - 9/7/21), CUDA 11.2.

[psiden]

I also get this error with TF 2.7