diff --git a/Transformer/Model.swift b/Transformer/Model.swift
index 70633007606..6d23851b16d 100644
--- a/Transformer/Model.swift
+++ b/Transformer/Model.swift
@@ -105,7 +105,7 @@ func _vjpCausallyMasked(_ dotProducts: Tensor<Float>, enable: Bool)
     return (causallyMasked(dotProducts), identity)
 }
 
-struct Attention: Layer {
+struct Attention: ParameterlessLayer {
     @noDerivative let dropout: Dropout<Float>
     @noDerivative let scale: Tensor<Float>
     @noDerivative let causal: Bool
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 801637fe2a0..fce54860149 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -3,7 +3,7 @@
 FROM nvidia/cuda:9.2-cudnn7-devel-ubuntu18.04
 
 # Allows the caller to specify the toolchain to use.
-ARG swift_tf_url=https://storage.googleapis.com/s4tf-kokoro-artifact-testing/latest/swift-tensorflow-DEVELOPMENT-ubuntu18.04.tar.gz
+ARG swift_tf_url=https://storage.googleapis.com/swift-tensorflow-artifacts/nightlies/latest/swift-tensorflow-DEVELOPMENT-ubuntu18.04.tar.gz
 
 # Install Swift deps.
 ENV DEBIAN_FRONTEND=noninteractive